DiffIR: Efficient Diffusion Model for Image Restoration

TL;DR

This work tackles the inefficiency of applying diffusion models to image restoration by introducing DiffIR, which uses a compact IR prior representation (IPR) to guide restoration. It consists of a compact IR prior extraction network (CPEN) and a Dynamic IRformer (DIRformer), trained in two stages: first to learn how to utilize an IPR, then to estimate the IPR from low-quality inputs via diffusion. The key contributions are the CPEN+DIRformer design with dynamic modulation, a diffusion process tailored to operate on a compact IPR, and a jointly trained pipeline that achieves state-of-the-art results in inpainting, super-resolution, and motion deblurring with substantially lower computational cost. Practically, DiffIR enables high-quality IR with far fewer diffusion steps and parameters, making diffusion-based IR more scalable and robust for real-world deployment.

Abstract

Diffusion model (DM) has achieved SOTA performance by modeling the image synthesis process into a sequential application of a denoising network. However, different from image synthesis, image restoration (IR) has a strong constraint to generate results in accordance with ground-truth. Thus, for IR, traditional DMs running massive iterations on a large model to estimate whole images or feature maps is inefficient. To address this issue, we propose an efficient DM for IR (DiffIR), which consists of a compact IR prior extraction network (CPEN), dynamic IR transformer (DIRformer), and denoising network. Specifically, DiffIR has two training stages: pretraining and training DM. In pretraining, we input ground-truth images into CPEN$_{S1}$ to capture a compact IR prior representation (IPR) to guide DIRformer. In the second stage, we train the DM to directly estimate the same IRP as pretrained CPEN$_{S1}$ only using LQ images. We observe that since the IPR is only a compact vector, DiffIR can use fewer iterations than traditional DM to obtain accurate estimations and generate more stable and realistic results. Since the iterations are few, our DiffIR can adopt a joint optimization of CPEN$_{S2}$, DIRformer, and denoising network, which can further reduce the estimation error influence. We conduct extensive experiments on several IR tasks and achieve SOTA performance while consuming less computational costs. Code is available at \url{https://github.com/Zj-BinXia/DiffIR}.

Figures (11)

• Figure 1: The Mult-Adds are measured on 256$\times$256 inputs. Our DiffIR achieves SOTA performance on IR tasks. Notably, LDM LDM and RePaint repaint are DM-based methods, and DiffIR is 1000$\times$ more efficient than RePaint while achieving better performance.
• Figure 2: The overview of the proposed DiffIR, which consists of DIRformer, CPEN, and denoising network. DiffIR has two training stages: (a) In the first stage, CPEN$_{S1}$ takes the ground-truth image as input and outputs an IPR $\mathbf{Z}$ to guide DIRformer to restore images. We optimize the CPEN$_{S1}$ with DiffIR$_{S1}$ together to make DiffIR$_{S1}$ can fully use extracted IPR. (b) In the second stage, we use the strong data estimation abilities of the DM to estimate the IPR extracted by pretrained CPEN$_{S1}$. Notably, we do not input the ground-truth image into CPEN$_{S2}$ and denoising networks. In the inference stage, we only use the reverse process of DM.
• Figure 3: Visual comparison of inpainting methods. Zoom-in for better details.
• Figure 4: Visual comparison of 4$\times$image super-resolution methods. Zoom-in for better details.
• Figure 5: Visual comparison of single image motion deblurring methods. Zoom-in for better details.