DiffIR: Efficient Diffusion Model for Image Restoration

In plain words

Image restoration means fixing a damaged picture: filling in a hole, making a blurry or tiny image sharp, or removing motion blur. Diffusion models are very good at generating detail, but they usually run many denoising steps over the whole image, which is slow and heavy. Much of a damaged picture is still there, so most of that work is wasted.

DiffIR changes what the diffusion model works on. A separate network looks at the clean image during training and squeezes the missing detail into a short vector of numbers. The diffusion model only has to guess that short vector from the damaged image, in a handful of steps. A transformer then uses the vector to rebuild the picture. Training happens in two stages, and the second stage trains all parts together so that small errors in the guess matter less.

TL;DR

• DiffIR diffuses a 256-value image-restoration prior (IPR) instead of pixels, with 4 diffusion steps §5.1 §4.1.
• It is trained in two stages: a pretrained encoder gives the prior from the ground-truth image, then a diffusion model predicts it from the low-quality image §4.1 §4.2.
• Author-reported state of the art on inpainting, x4 super-resolution and motion deblurring §5.2 §5.3 §5.4.
• Official code is Apache-2.0 at commit 293f86c, but several details in the code differ from the paper (see Paper vs code).

Problem

Diffusion models for restoration usually reuse the image-synthesis recipe: many iterations, a large network, and denoising over whole images or feature maps. The authors argue that in restoration most pixels are given, so this is wasteful, and that it can create details that do not match the input §1 §4. The diffusion restorers they compare with are RePaint, LDM and SRdiff §2 tab. 1 tab. 2.

Method

DiffIR has three parts: a compact IR prior extraction network (CPEN), a dynamic IR transformer (DIRformer), and a denoising network §4.

Stage 1, pretraining §4.1. The ground-truth image and the low-quality (LQ) image are combined, unshuffled and encoded by CPEN_S1 into the IPR Z with 4C' values eq. 5. The IPR modulates the DIRformer's feed-forward (DGFN) and attention (DMTA) blocks through linear maps and element-wise products eq. 6 eq. 7 eq. 8. CPEN_S1 and DIRformer are trained together with an L1 reconstruction loss eq. 9.

Stage 2, diffusion on the IPR §4.2. The diffusion process is applied to the stage-1 IPR Z, not to pixels: Z_T is drawn from q(Z_T | Z) eq. 10. The reverse process runs all T steps and produces an estimate of Z. The paper writes the reverse update without a variance term eq. 11. A conditioning vector D is computed from the LQ image by CPEN_S2 eq. 12. The denoising network takes the current estimate, the step and D.

Joint training §4.2 eq. 13. The IPR estimate is matched to the stage-1 IPR with an L1 loss L_diff. It is added to the reconstruction loss L_rec, and the DIRformer is fed the diffused estimate, so that CPEN_S2, the denoiser and DIRformer are optimised together. The paper says perceptual and adversarial losses can be added for visual quality §4.2.

Inference §4.2 alg. 1 alg. 2. Draw a Gaussian noise vector, run the reverse process for T steps conditioned on D, then pass the resulting IPR to DIRformer, which restores the LQ image. No ground truth is needed.

Pseudocode, written from the paper. Equation numbers refer to the paper.

# Stage 1: prior from ground truth  [§4.1] [eq. 5]
# (B, 3, H, W) images; the IPR has 4C' values, C' = 64  [§5.1]
def stage1_ipr(lq_image, gt_image):
    # Concatenate, unshuffle, encode. Read literally, see Pitfalls.
    ipr_z = cpen_s1(pixel_unshuffle(concat(gt_image, lq_image)))
    return ipr_z                                # (B, 4C')

# Stage 1 loss  [eq. 9]
def stage1_loss(lq_image, gt_image):
    ipr_z = stage1_ipr(lq_image, gt_image)
    # DIRformer uses the IPR to modulate DGFN and DMTA  [eq. 6-8]
    restored = dirformer(lq_image, ipr_z)
    return l1_mean(gt_image - restored)

# One reverse step, eq. 11 as written (no variance term)  [§4.2]
def reverse_step(ipr_t, t, condition):
    # Denoiser output used as the noise estimate  [eq. 11]
    # The code predicts the clean IPR instead; see Pitfalls.
    predicted_noise = denoiser(concat(ipr_t, t, condition))
    shifted = ipr_t - predicted_noise * (1 - alpha_t) / sqrt(1 - alpha_bar_t)
    return shifted / sqrt(alpha_t)              # [eq. 11]

# Stage 2: joint training step  [§4.2] [eq. 10] [eq. 12] [eq. 13]
def stage2_loss(lq_image, gt_image, ipr_z):
    # Noise the stage-1 IPR to step T  [eq. 10]
    noise = random_normal_like(ipr_z)
    ipr_t = sqrt(alpha_bar_T) * ipr_z + sqrt(1 - alpha_bar_T) * noise
    # Condition vector from the LQ image only  [eq. 12]
    condition = cpen_s2(lq_image)               # (B, 4C')
    # Run all T reverse steps, keeping gradients  [§4.2]
    for t in reversed(range(T)):
        ipr_t = reverse_step(ipr_t, t, condition)
    ipr_hat = ipr_t
    restored = dirformer(lq_image, ipr_hat)
    # L_diff and L_all, eq. 13 (L1 means over the 4C' values)
    l_diff = l1_mean(ipr_hat - ipr_z)
    return l1_mean(gt_image - restored) + l_diff

# Inference  [§4.2] [alg. 2]
def restore(lq_image):
    condition = cpen_s2(lq_image)               # (B, 4C')
    # Start from N(0, 1), not from q(Z_T | Z)  [alg. 2]
    ipr_t = random_normal((batch, 4 * C_prime))
    for t in reversed(range(T)):
        ipr_t = reverse_step(ipr_t, t, condition)
    return dirformer(lq_image, ipr_t)

Key equations as stated in the paper:

q(x_t \mid x_0) = \mathcal{N}(x_t;\ \sqrt{\bar\alpha_t}\,x_0,\ (1-\bar\alpha_t)\,I) eq. 2

\hat{Z}_{t-1} = \tfrac{1}{\sqrt{\alpha_t}}\big(\hat{Z}_t - \epsilon\, \tfrac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\big) eq. 11

\mathcal{L}_{diff} = \tfrac{1}{4C'}\textstyle\sum_{i=1}^{4C'} |\hat{Z}(i) - Z(i)|, \quad \mathcal{L}_{all} = \mathcal{L}_{rec} + \mathcal{L}_{diff} eq. 13

The paper's ablation also compares an L2 loss eq. 14 and a KL loss eq. 15 on the IPR.

Hyperparameters

Values in the table come from the paper unless the Notes say otherwise.

Results

Author-reported. Each table is the paper's own. Setup is given per table. Best and second-best marks (bold and underline in the source) are not reproduced here. Read the values as shown.

Inpainting tab. 1. FID and LPIPS, lower is better. Places at 512 x 512, CelebA-HQ at 256 x 256. Narrow and wide masks. Method size in M parameters.

Super-resolution, x4 tab. 2. PSNR (higher) and LPIPS (lower) on five sets. The caption names LPIPS and DISTS, but the columns are PSNR and LPIPS. Only GAN-based methods and the diffusion methods LDM and SRdiff are compared.

Motion deblurring tab. 3. PSNR and SSIM, higher is better. Trained on GoPro.

Ablations on CelebA-HQ inpainting FID tab. 4 tab. 5 §6. All rows use the same LQ input at 256 x 256. Mult-Adds are for that size.

Real-world 4x SR tab. 6, appendix app. B. LPIPS, DISTS (lower is better), PSNR (higher is better). Mult-Adds in T, at LR 256 x 256.

Headline numbers in the text, with the paper's setup:

• Inpainting: DiffIR_S2 has up to 0.2706 and 0.5583 lower wide-mask FID than LaMa on Places and CelebA-HQ, at similar parameters and Mult-Adds §5.2.
• SR: over BebyGAN, the LPIPS margin is up to 0.0151 (DIV2K100) and 0.0089 (Urban100), at 63% of the compute §5.3.
• Deblurring: 0.28 dB over Restormer on GoPro and 0.33 dB on HIDE, at 78% of Restormer's Mult-Adds §5.4.
• Efficiency: "1000x more efficient than RePaint" is shown in fig. 1 and is not derivable from the tables §1.

Paper vs code

Code paths are relative to the repository root at code_commit (raw/code/2303.09472/). The abstract links the repo, and the code is Apache-2.0. The repo has four variants: DiffIR-SRGAN (SR), DiffIR-inpainting, DiffIR-demotionblur, and DiffIR-RealSR. The SR variant was read most closely. The deblurring configs were checked for hyperparameters. Inpainting was checked for architecture and losses. RealSR was only checked for its S2 config values.

1. Beta schedule. The paper says beta rises linearly from 0.1 to 0.99 §5.1. The code interpolates sqrt(beta) linearly and squares it util2.py. At T = 4 the code gives about [0.1, 0.294, 0.591, 0.99]. Linear in beta would give [0.1, 0.397, 0.693, 0.99].
2. What the denoiser predicts. Eq. 11 uses a noise estimate epsilon eq. 11 §4.2. The code builds the diffusion module with the default parameterization="x0" ddpm.py. The network output is treated as the clean IPR estimate, and the step uses the DDPM posterior mean computed from it ddpm.py.
3. Timestep indexing and time embedding. The paper's steps run 1..T. The code runs 0..T-1 ddpm.py. The denoiser's time input is t / (10 T), not a sinusoidal embedding S2_arch.py.
4. Denoiser architecture. The paper does not specify the denoiser. The code uses an MLP: one linear layer over the concatenation of condition (256), time (1) and IPR (256), then n_denoise_res linear + LeakyReLU blocks of width 256, without a residual connection S2_arch.py.
5. Attention normalisation. The paper scales Q K^T by a learnable gamma eq. 7. The code L2-normalises Q and K over the token dimension before the product, then multiplies by a learnable per-head temperature S2_arch.py.
6. Extra KD loss in stage 2. The paper's total loss is L_rec + L_diff eq. 13. The code adds a KL term between softmax-normalised IPRs (temperature 0.15), on top of the L1 IPR term and the pixel L1 loss DiffIR_S2_model.py my_loss.py. The KL weight is 1 and the temperature 0.15 in the SR config train_DiffIRS2_x4.yml.
7. Stage-1 input for x4 SR. Eq. 5 unshuffles the concatenation of GT and LQ. For x4 SR the code takes the raw LQ image (3 channels) and concatenates it with the 4x-unshuffled GT (48 channels), giving 51 input channels S1_arch.py S1_arch.py. For inpainting and x2, each image is unshuffled before concatenation.
8. Stage-2 CPEN input for x4 SR. Eq. 12 unshuffles the LQ image. The SR code does not, and feeds the raw 3-channel LQ into CPEN_S2 S2_arch.py.
9. Channel widths. The paper gives [48, 96, 192, 384] §5.1. The SR and RealSR configs use dim 64 train_DiffIRS2_x4.yml. Inpainting uses dim 52 DiffIRS2-place2.yaml. Deblurring uses dim 48, as the paper says train_DiffIRS2.yml.
10. Deblurring schedule. The paper gives stage-1 patch and batch pairs and iterations at 92K, 156K, 204K, 240K and 276K app. F. The code's stage-1 progressive patch sizes run 192 to 512, with iterations [92000, 80000, 38000, 33000, 33000, 24000] (cumulative about 92K, 172K, 210K, 243K, 276K, 300K) train_DiffIRS1.yml. The 156K, 204K and 240K points do not match the code.
11. Deblurring stage-2 warm-up. The code has an encoder-only warm-up of encoder_iter: 100000 steps, then a joint phase up to 400K steps with MultiStepLR train_DiffIRS2.yml. During warm-up the step's loss is only the IPR L1 term DiffIR_S2_model.py. This is the SR model's training step; the deblurring copy of it was not diffed line by line. The paper does not mention this phase, and gives 300K iterations with cosine decay app. F.
12. Learning-rate schedule and EMA. The paper states no schedule for SR or inpainting. The SR config uses MultiStepLR with milestones 500K and 700K and gamma 0.5 train_DiffIRS2_x4.yml, and EMA 0.999 train_DiffIRS2_x4.yml.
13. Inpainting batch. The paper says batch 30 §5.2. The launch script uses data.batch_size=4 on 8 GPUs, so 32 in total train_place256S2.sh. The CelebA data config uses batch_size: 5 abl-04-256-mh-dist-celeba.yaml.
14. Stage-1 frozen. The code puts the stage-1 encoder in eval mode and excludes it from the stage-2 optimiser DiffIR_S2_model.py. The paper does not say whether CPEN_S1 is frozen.
15. Real-world SR loss. The appendix says it uses the losses of "Real-ESRGAN~\cite{ESRGAN}" app. B. The citation key points to ESRGAN. The paper may mean Real-ESRGAN. Not resolved.

No discrepancies were checked for the RealSR and inpainting denoisers beyond the configs listed above.

Pitfalls

• Beta schedule. A literal linear beta schedule, as the paper states, is not what the code trains. Use the sqrt-linear schedule if you want to reproduce the released numbers. Paper vs code, item 1.
• Noise or clean-signal prediction. If you implement eq. 11 with a noise head, you will not match the code's x0 head, and the two have different training targets. Paper vs code, item 2.
• Index convention. Steps are 0-based in the code. The time scaling divides by 10 T, so t = 3 enters as 0.075 when T = 4. Paper vs code, item 3.
• Starting noise. Training starts from q(Z_T | Z), inference from N(0, 1). With the code's schedule, the product of alphas is about 0.0026, so sqrt(alpha_bar_T) is about 0.05 and the start is close to pure noise. This is computed from the code's betas, not stated in the paper.
• Eq. 5 and eq. 12 read literally do not match the x4 SR code (see Paper vs code, items 7 and 8). For inpainting and x2 the code unshuffles each image before concatenation, which is different again from eq. 5, which unshuffles the concatenation. Check the shapes in your own implementation.
• Attention. The L2 normalisation of Q and K is not in the paper. An implementation that follows only the paper will differ from the released code. Paper vs code, item 5.
• Undocumented loss terms. The KL term in stage 2 is not in the paper and changes the objective. Paper vs code, item 6.
• Patch size ambiguity. For SR the main text says the LQ patch is 64 x 64 §5.3. The appendix says random crops of 256 x 256 app. E. The code's gt_size is 256 (the HR side), which matches the LQ patch of 64 train_DiffIRS2_x4.yml. Use 64 LQ, 256 HR.
• Deblurring patch literal. The appendix writes the patch pairs as "(1602, 40)" etc. The intended size is most likely 160 x 160 (the same pattern as 128 x 128 in the text). The code's sizes start at 192. Verify before copying.
• Compute claims. "1000x more efficient than RePaint" fig. 1 and "0.1% computational resources" §5.2 are not in any table. Mult-Adds for the inpainting, SR and deblurring comparisons are not in the tables either; they appear in fig. 1 and in the text. Cite them as author-reported.
• Text and table mismatches.
  • §5.2: "compared with DM based method RePaint~\cite{LDM}". The citation key points to LDM.
  • §5.4: "surpasses IPT and MIMI-Unet+ by 0.68 dB and 0.54 dB on GoPro". The table gives MIMO-Unet+ 32.45, so the margin is 0.75 dB. The 0.54 dB matches MPRNet (32.66).
  • §5.3 mentions Set5 and lists it as a test set. Table 2 has no Set5 column.
  • tab. 2 caption says LPIPS/DISTS; the columns are PSNR and LPIPS.
  • §6 says the variant with joint training has "quite similar LPIPS" to stage 1. tab. 4 reports FID.
• Missing details. No learning-rate schedule for SR and inpainting, no loss weights for the perceptual and adversarial terms, no denoiser architecture, and no inference timing.
• Inference passes. Each inference runs T = 4 denoiser passes and one DIRformer pass ddpm.py. The paper's memory and latency numbers for this are not given beyond Mult-Adds.
• Appendix letters. Appendix citations (app. B to app. F) follow the section order in the LaTeX source. The PDF may label them differently.

Lineage

• Builds on: the DDPM denoising diffusion formulation, as the paper's preliminaries state §3. The DIRformer blocks (DMTA and DGFN) are the paper's own design §4.1.
• Compared against: GAN restorers (SFTGAN, SRGAN, ESRGAN, USRGAN, SPSR, BebyGAN, KDSR, BSRGAN, Real-ESRGAN), non-diffusion restorers (EdgeConnect, ICT, LaMa, MT-RNN, MPRNet, Restormer, MIMO-Unet+, IPT and others), and diffusion restorers (LDM, RePaint, SRdiff) tab. 1 tab. 2 tab. 3 tab. 6.
• Superseded by: not known from content already in this wiki.