# Selective Hourglass Mapping for Universal Image Restoration Based on Diffusion Model

arXiv: https://arxiv.org/abs/2403.11157 · HTML: https://mlsanity.com/paper/2403.11157/

## Key facts

- Cited by 73 papers in the map (data through 14 Oct 2026).
- Placed in Inverse problems and image restoration, one of 14 directions in Diffusion models.
- Builds on 56 of its 72 references that are in the map.
- First posted on arXiv 17 Mar 2024.

## Placed in

- Diffusion models / [Inverse problems and image restoration](https://mlsanity.com/diffusion/inverse-problems-and-image-restoration/)

## Builds on (most cited first, 30 of 56)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [SwinIR: Image Restoration Using Swin Transformer](https://mlsanity.com/paper/2108.10257/) (2021)
- [Palette: Image-to-Image Diffusion Models](https://mlsanity.com/paper/2111.05826/) (2021)
- [An Introduction to Variational Autoencoders](https://mlsanity.com/paper/1906.02691/) (2019)
- [Blended Diffusion for Text-driven Editing of Natural Images](https://mlsanity.com/paper/2111.14818/) (2021)
- [Multi-Stage Progressive Image Restoration](https://mlsanity.com/paper/2102.02808/) (2021)
- [Deep Multi-scale Convolutional Neural Network for Dynamic Scene Deblurring](https://mlsanity.com/paper/1612.02177/) (2016)
- [Deep Retinex Decomposition for Low-Light Enhancement](https://mlsanity.com/paper/1808.04560/) (2018)
- [Simple Baselines for Image Restoration](https://mlsanity.com/paper/2204.04676/) (2022)
- [DiffEdit: Diffusion-based semantic image editing with mask guidance](https://mlsanity.com/paper/2210.11427/) (2022)
- [Zero-Shot Image Restoration Using Denoising Diffusion Null-Space Model](https://mlsanity.com/paper/2212.00490/) (2022)
- [Label-Efficient Semantic Segmentation with Diffusion Models](https://mlsanity.com/paper/2112.03126/) (2021)
- [Benchmarking Single Image Dehazing and Beyond](https://mlsanity.com/paper/1712.04143/) (2017)
- [Open-Vocabulary Panoptic Segmentation with Text-to-Image Diffusion Models](https://mlsanity.com/paper/2303.04803/) (2023)
- [DiffusionDet: Diffusion Model for Object Detection](https://mlsanity.com/paper/2211.09788/) (2022)
- [Deep Joint Rain Detection and Removal from a Single Image](https://mlsanity.com/paper/1609.07769/) (2016)
- [Images Speak in Images: A Generalist Painter for In-Context Visual Learning](https://mlsanity.com/paper/2212.02499/) (2022)
- [MAXIM: Multi-Axis MLP for Image Processing](https://mlsanity.com/paper/2201.02973/) (2022)
- [FreeDoM: Training-Free Energy-Guided Conditional Diffusion Model](https://mlsanity.com/paper/2303.09833/) (2023)
- [Multi-Scale Progressive Fusion Network for Single Image Deraining](https://mlsanity.com/paper/2003.10985/) (2020)
- [Vision Transformers for Single Image Dehazing](https://mlsanity.com/paper/2204.03883/) (2022)

## Built on by (most cited first, 30 of 73)

- [Diffusion Models in Low-Level Vision: A Survey](https://mlsanity.com/paper/2406.11138/) (2024)
- [A Survey on All-in-One Image Restoration: Taxonomy, Evaluation and Future Trends](https://mlsanity.com/paper/2410.15067/) (2024)
- [UniRestore: Unified Perceptual and Task-Oriented Image Restoration Model Using Diffusion Prior](https://mlsanity.com/paper/2501.13134/) (2025)
- [FoundIR: Unleashing Million-scale Training Data to Advance Foundation Models for Image Restoration](https://mlsanity.com/paper/2412.01427/) (2024)
- [Conditional Image Synthesis with Diffusion Models: A Survey](https://mlsanity.com/paper/2409.19365/) (2024)
- [Decoupled Distillation to Erase: A General Unlearning Method for Any Class-centric Tasks](https://mlsanity.com/paper/2503.23751/) (2025)
- [Visual-Instructed Degradation Diffusion for All-in-One Image Restoration](https://mlsanity.com/paper/2506.16960/) (2025)
- [Mitigating the Human-Robot Domain Discrepancy in Visual Pre-training for Robotic Manipulation](https://mlsanity.com/paper/2406.14235/) (2024)
- [UniUIR: Considering Underwater Image Restoration as An All-in-One Learner](https://mlsanity.com/paper/2501.12981/) (2025)
- [WeatherBench: A Real-World Benchmark Dataset for All-in-One Adverse Weather Image Restoration](https://mlsanity.com/paper/2509.11642/) (2025)
- [AWRaCLe: All-Weather Image Restoration using Visual In-Context Learning](https://mlsanity.com/paper/2409.00263/) (2024)
- [LoRA-IR: Taming Low-Rank Experts for Efficient All-in-One Image Restoration](https://mlsanity.com/paper/2410.15385/) (2024)
- [RestoreVAR: Visual Autoregressive Generation for All-in-One Image Restoration](https://mlsanity.com/paper/2505.18047/) (2025)
- [Residual Diffusion Bridge Model for Image Restoration](https://mlsanity.com/paper/2510.23116/) (2025)
- [Learning Continuous Wasserstein Barycenter Space for Generalized All-in-One Image Restoration](https://mlsanity.com/paper/2602.23169/) (2026)
- [GenDeg: Diffusion-based Degradation Synthesis for Generalizable All-In-One Image Restoration](https://mlsanity.com/paper/2411.17687/) (2024)
- [Modeling Multiple Normal Action Representations for Error Detection in Procedural Tasks](https://mlsanity.com/paper/2503.22405/) (2025)
- [Panorama Generation From NFoV Image Done Right](https://mlsanity.com/paper/2503.18420/) (2025)
- [IRBridge: Solving Image Restoration Bridge with Pre-trained Generative Diffusion Models](https://mlsanity.com/paper/2505.24406/) (2025)
- [FoundIR-v2: Optimizing Pre-Training Data Mixtures for Image Restoration Foundation Model](https://mlsanity.com/paper/2512.09282/) (2025)
- [Efficient Degradation-agnostic Image Restoration via Channel-Wise Functional Decomposition and Manifold Regularization](https://mlsanity.com/paper/2505.18679/) (2025)
- [DSwinIR: Rethinking Window-based Attention for Image Restoration](https://mlsanity.com/paper/2504.04869/) (2025)
- [M2Restore: Mixture-of-Experts-based Mamba-CNN Fusion Framework for All-in-One Image Restoration](https://mlsanity.com/paper/2506.07814/) (2025)
- [UniCoRN: Latent Diffusion-based Unified Controllable Image Restoration Network across Multiple Degradations](https://mlsanity.com/paper/2503.15868/) (2025)
- [CoA: Towards Real Image Dehazing via Compression-and-Adaptation](https://mlsanity.com/paper/2504.05590/) (2025)
- [Beyond Degradation Redundancy: Contrastive Prompt Learning for All-in-One Image Restoration](https://mlsanity.com/paper/2504.09973/) (2025)
- [Bilevel Layer-Positioning LoRA for Real Image Dehazing](https://mlsanity.com/paper/2603.10872/) (2026)
- [Learning Efficient and Effective Trajectories for Differential Equation-based Image Restoration](https://mlsanity.com/paper/2410.04811/) (2024)
- [DGSolver: Diffusion Generalist Solver with Universal Posterior Sampling for Image Restoration](https://mlsanity.com/paper/2504.21487/) (2025)
- [Enhancing Image Restoration Transformer via Adaptive Translation Equivariance](https://mlsanity.com/paper/2506.18520/) (2025)

## Abstract (arXiv)

Universal image restoration is a practical and potential computer vision task for real-world applications. The main challenge of this task is handling the different degradation distributions at once. Existing methods mainly utilize task-specific conditions (e.g., prompt) to guide the model to learn different distributions separately, named multi-partite mapping. However, it is not suitable for universal model learning as it ignores the shared information between different tasks. In this work, we propose an advanced selective hourglass mapping strategy based on diffusion model, termed DiffUIR. Two novel considerations make our DiffUIR non-trivial. Firstly, we equip the model with strong condition guidance to obtain accurate generation direction of diffusion model (selective). More importantly, DiffUIR integrates a flexible shared distribution term (SDT) into the diffusion algorithm elegantly and naturally, which gradually maps different distributions into a shared one. In the reverse process, combined with SDT and strong condition guidance, DiffUIR iteratively guides the shared distribution to the task-specific distribution with high image quality (hourglass). Without bells and whistles, by only modifying the mapping strategy, we achieve state-of-the-art performance on five image restoration tasks, 22 benchmarks in the universal setting and zero-shot generalization setting. Surprisingly, by only using a lightweight model (only 0.89M), we could achieve outstanding performance. The source code and pre-trained models are available at https://github.com/iSEE-Laboratory/DiffUIR