Selective Hourglass Mapping for Universal Image Restoration Based on Diffusion Model

TL;DR

The paper tackles universal image restoration by replacing multi-task, distribution-specific learning with a selective hourglass diffusion framework (DiffUIR) that unifies degradation distributions via a shared distribution term (SDT) while maintaining strong condition guidance. By injecting degradation conditions into the forward diffusion and gradually reducing their influence, the model learns a common distribution from which task-specific outputs are recovered through a guided reverse process. The approach yields state-of-the-art results across five restoration tasks (deraining, low-light, desnowing, dehazing, deblurring) on 22 benchmarks, including zero-shot real-world scenarios, with a lightweight model DiffUIR-T at 0.89M parameters. The method demonstrates the practical potential of universal restoration with significantly reduced parameter and compute footprints, supported by ablations and distribution analyses that highlight the benefits of SDT and explicit conditioning.

Abstract

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

Figures (7)

• Figure 1: An illustration of existing universal image restoration methods compared with our DiffUIR, existing methods mainly design task-specific modules to handle different distributions, which force the generic model (tangerine module) to learn different distributions at once, termed multi-partite mapping. In contrast, the proposed DiffUIR maps the different distributions to one shared distribution (i.e., note that it is not the pure Gaussian distribution) while maintaining strong condition guidance. In this way, DiffUIR enables the generic model to only learn one shared distribution and guides the shared distribution to a task-specific distribution, termed selective hourglass mapping. Zoom in for best view.
• Figure 2: Visualization comparison with state-of-the-art methods on low-light enhancement. Zoom in for best view.
• Figure 3: Visualization comparison with state-of-the-art methods on deraining. Zoom in for best view.
• Figure 4: Visualization comparison with state-of-the-art methods on debluring. Zoom in for best view.
• Figure 5: Distributions of the diffusing endpoint features of different image restoration tasks generated by w/o SDT and with SDT (value 0.9) visualized by t-SNE tsne. The diffusing endpoint features of different tasks are in different distributions without any interaction when without SDT. With the SDT, DiffUIR could achieve shared distribution mapping and strong condition guidance simultaneously. Zoom in for best view.