Refusion: Enabling Large-Size Realistic Image Restoration with Latent-Space Diffusion Models

Ziwei Luo; Fredrik K. Gustafsson; Zheng Zhao; Jens Sjölund; Thomas B. Schön

Source-linked AI summary

Refusion: Enabling Large-Size Realistic Image Restoration with Latent-Space Diffusion Models

Ziwei Luo, Fredrik K. Gustafsson, Zheng Zhao, Jens Sjölund, Thomas B. Schön

arXiv:2304.08291v1Apr 17, 2023cs.CV

TL;DR

Realistic image restoration remains difficult for diffusion models because real-world degradations and image sizes are complex, while conventional restoration methods can produce over-smooth images. Refusion improves diffusion architecture and training choices and introduces U-Net latent diffusion for large images. It performs effectively across several restoration tasks, including 6000 × 4000 × 3 inputs, and achieves the best perceptual performance in the NTIRE 2023 Shadow Removal Challenge.

Problem
Image restoration must recover high-quality images from degraded inputs, but diffusion models remain challenged by complex real-world degradations and large, arbitrary image sizes.
Method
Refusion combines IR-SDE-based diffusion improvements with NAFNet noise prediction and U-Net latent diffusion that preserves high-resolution input information for decoding.
Results
Refusion is effective across shadow removal, high-resolution dehazing, stereo super-resolution, and bokeh transformation, and achieves the best perceptual performance in NTIRE 2023 Shadow Removal.
Takeaways & Limitations
U-Net latent Refusion extends diffusion-based restoration to large images, including 6000 × 4000 × 3 pixels, while recovering accurate details.

Abstract

from arXiv · show

This work aims to improve the applicability of diffusion models in realistic image restoration. Specifically, we enhance the diffusion model in several aspects such as network architecture, noise level, denoising steps, training image size, and optimizer/scheduler. We show that tuning these hyperparameters allows us to achieve better performance on both distortion and perceptual scores. We also propose a U-Net based latent diffusion model which performs diffusion in a low-resolution latent space while preserving high-resolution information from the original input for the decoding process. Compared to the previous latent-diffusion model which trains a VAE-GAN to compress the image, our proposed U-Net compression strategy is significantly more stable and can recover highly accurate images without relying on adversarial optimization. Importantly, these modifications allow us to apply diffusion models to various image restoration tasks, including real-world shadow removal, HR non-homogeneous dehazing, stereo super-resolution, and bokeh effect transformation. By simply replacing the datasets and slightly changing the noise network, our model, named Refusion, is able to deal with large-size images (e.g., 6000 x 4000 x 3 in HR dehazing) and produces good results on all the above restoration problems. Our Refusion achieves the best perceptual performance in the NTIRE 2023 Image Shadow Removal Challenge and wins 2nd place overall.

1. Introduction

Refusion improves diffusion-based realistic image restoration by addressing complex degradations, arbitrary image sizes, and over-smoothing, while supporting diverse restoration tasks and high-resolution inputs.

Motivation: Diffusion restoration remains challenging for complex degradations and the large, arbitrary image sizes found in real-world datasets.Existing pretrained approaches may also require carefully curated datasets or known degradation parameters.
Refusion approach: Refusion builds on IR-SDE and studies noise-network architecture, noise levels, denoising steps, training image sizes, and optimizer/scheduler choices.The study targets improved diffusion effectiveness across diverse real-world restoration problems.
Applications: Refusion is effective on real-world shadow removal, high-resolution non-homogeneous dehazing, stereo super-resolution, and bokeh effect transformation.The approach was evaluated on extensive real-world and synthetic datasets.
High-resolution restoration: U-Net latent compression enables high-resolution restoration while preserving input information for decoding, improving stability and recovering accurate images without adversarial optimization.Diffusion is performed in a compressed latent space, while high-resolution information from the original input is retained.
Refusion approach: NAFNet achieves better image restoration performance across all tasks while using fewer parameters and greater computational efficiency than the diffusion U-Net base network.The authors replace the diffusion base network with NAFNet.

2. Related Work

Related work spans CNN, transformer, and nonlinear-activation-free architectures for image restoration, while Figure 2 presents Refusion’s task and NAFBlock designs.

CNN-based restoration: SRCNN and DnCNN helped establish CNN-based approaches for image super-resolution and denoising, stimulating broader image-restoration research.These early contributions significantly improved performance in their respective applications.
Transformer-based restoration: Transformer-based architectures have been applied to image restoration, with IPT, SwinIR, Restormer, and Uformer targeting efficient or high-performing restoration designs.These approaches combine transformer components with convolutional, residual, or U-shaped structures.
Refusion design: Figure 2(a) depicts Refusion across restoration tasks, while Figure 2(b) shows a NAFBlock with an added time-processing branch.SimpleGate splits feature channels into two parts and multiplies them element-wise; SCA denotes simple channel attention.
Alternative architectures: Other restoration efforts use MLPs and nonlinear-activation-free networks as alternatives to conventional activation-based architectures.Refusion’s NAFNet-based design belongs to this latter line of work.

3. Preliminaries: Mean-Reverting SDE

Refusion uses IR-SDE, a mean-reverting diffusion framework that transforms high-quality images toward degraded counterparts and reverses this process for restoration. Its forward SDE approaches mean μ with fixed noise level λ, while a CNN estimates noise for reverse-time sampling.

IR-SDE models degradation by transforming high-quality images into low-quality counterparts through a forward-time Itô SDE.The reverse-time process then transforms low-quality images toward high-quality versions.
The forward process uses drift θ_t(μ − x) and stochastic volatility σ_t, controlling mean reversion and random perturbation.θ_t and σ_t are time-dependent positive parameters.
Under σ_t^2/θ_t = 2λ^2, the marginal distribution is Gaussian with mean m_t and variance v_t.The mean and variance are defined from the accumulated mean-reversion rate.
As t increases, m_t converges to μ and v_t converges to λ^2, so the initial state is iteratively driven toward μ with fixed noise level λ.This establishes the stationary behavior of the forward process.
During training, the ground-truth score is computed from available high-quality images, while a CNN learns to estimate the noise for reverse sampling.The reparameterization expresses x(t) as m_t(x) plus scaled standard Gaussian noise.

4. Improving the Diffusion Model

Refusion improves diffusion-based restoration through latent-space processing, a more efficient NAFNet noise predictor, and empirical tuning of training strategies. The method targets high-resolution and diverse restoration settings while improving training stability and performance.

4.1. Latent Diffusion under U-Net: The U-Net latent model performs restoration in a low-resolution space, enabling diffusion on large images while preserving high-resolution input information for decoding.This design addresses the cost of iterative diffusion on high-resolution inputs, including 6000 × 4000 × 3 dehazing images.
4.1. Latent Diffusion under U-Net: Latent-replacing pretraining uses two reconstruction L1 losses: one reconstructs the low-quality input, and another reconstructs the corresponding high-quality image after latent replacement.The strategy trains the decoder to recover high-quality images from transformed low-quality latent representations.
4.2. Modified NAFBlocks for Noise Prediction: NAFNet replaces U-Net as the diffusion base network, using nonlinear activation-free blocks and requiring fewer parameters with greater computational efficiency.Task-specific modifications include lens information for bokeh transformation and dual inputs for stereo super-resolution.
4.2. Modified NAFBlocks for Noise Prediction: The modified NAFNet significantly outperforms the U-Net backbone on the shadow removal task.Figure 6 compares their learning curves on the shadow removal dataset.
4.3. Improved Training Strategies: Noise level, denoising steps, and optimizer/scheduler choices materially affect diffusion restoration training on real-world shadow removal.Noise levels σ = 50 or σ = 70 are more stable than σ = 10, while the supplied passage introduces comparisons for denoising steps and optimization schedules.
4.3. Improved Training Strategies: Larger training patches perform much better than 128 × 128 patches in diffusion-based restoration.The comparison uses 128 × 128 and 256 × 256 training patches on the shadow removal task.

5. Experiments

Refusion is evaluated across four restoration tasks spanning shadow removal, stereo super-resolution, bokeh transformation, and extremely high-resolution non-homogeneous dehazing. Across these experiments, it improves perceptual restoration, supports large images through latent diffusion, and offers competitive runtime or efficiency advantages.

Tasks and datasets: Refusion is evaluated on shadow removal, stereo super-resolution, bokeh effect transformation, and HR non-homogeneous dehazing.The tasks include both synthetic and real-world restoration settings.
Implementation details: Training uses 256 × 256 patches for standard diffusion tasks and 1024 × 1024 patches compressed to 128 × 128 for latent diffusion tasks.Latent tasks first train the U-Net for 300,000 iterations, then train Refusion for 400,000 iterations.
Stereo image super-resolution: Refusion achieves significantly better performance than IR-SDE across all stereo super-resolution metrics and outperforms NAFSSR on perceptual scores.NAFSSR retains the best PSNR and SSIM, while Refusion improves LPIPS and FID relative to both baselines.
Evaluation: The experiments use LPIPS and FID as primary perceptual metrics, with PSNR and SSIM reported for reference and runtime compared across methods.Shadow removal additionally reports RMSE.
Image shadow removal: Refusion achieves the best overall restoration performance on shadow removal and runs significantly faster than IR-SDE.The comparison includes IR-SDE, DHAN, and an L1-trained U-Net baseline.
Bokeh effect transformation: Latent Refusion improves bokeh transformation perceptual performance over Restormer and reduces computation flops by about 15× compared with the original Refusion model.The latent variant achieves better LPIPS but slightly worse FID than Refusion, while running close to Restormer's runtime.
HR non-homogeneous dehazing: For 6000 × 4000 × 3 dehazing images, latent Refusion can process the task in a low-resolution latent space, whereas IR-SDE and Refusion cannot process a complete image at test time.Most fog is successfully removed in the reported visual results.

6. Conclusion

The paper concludes that Refusion improves the applicability of diffusion-based image restoration across diverse tasks. Its U-Net-based latent model performs diffusion in a low-resolution space while retaining input information for detailed reconstruction, including at 6000 × 4000 × 3 resolution.

Conclusion: Refusion is successfully applied to various image restoration tasks and achieves the best perceptual performance in the NTIRE 2023 Shadow Removal Challenge.The paper presents these results as evidence of improved diffusion-based restoration applicability.
Latent Refusion: The latent Refusion model compresses inputs into a low-resolution latent space for diffusion while preserving high-resolution information through decoder-connected hidden vectors.This design is reported to recover more accurate image details.
Large-size restoration: Latent Refusion can run on images sized 6000 × 4000 × 3 pixels.The capability is stated as part of the proposed approach for processing large-size images.