Denoising Diffusion Models for Plug-and-Play Image Restoration

Yuanzhi Zhu, Kai Zhang, Jingyun Liang, Jiezhang Cao, Bihan Wen, Radu Timofte, Luc Van Gool.

[project page] [paper]

This repository contains the code and data associated with the paper "Denoising Diffusion Models for Plug-and-Play Image Restoration", which was presented at the CVPR workshop NTIRE 2023.

This code is based on the OpenAI Guided Diffusion and DPIR.

Contents

• Denoising Diffusion Models for Plug-and-Play Image Restoration
  • Abstract
  • Setting Up
  • Clone and Install
  • Model Download
  • Inference Code
  • Brief Introduction
  • Results
  • Quantitative Results
  • Qualitative Results
  • Citation
  • Acknowledgments

Abstract

Plug-and-play Image Restoration (IR) has been widely recognized as a flexible and interpretable method for solving various inverse problems by utilizing any off-the-shelf denoiser as the implicit image prior. However, most existing methods focus on discriminative Gaussian denoisers. Although diffusion models have shown impressive performance for high-quality image synthesis, their potential to serve as a generative denoiser prior to the plug-and-play IR methods remains to be further explored. While several other attempts have been made to adopt diffusion models for image restoration, they either fail to achieve satisfactory results or typically require an unacceptable number of Neural Function Evaluations (NFEs) during inference. This paper proposes DiffPIR, which integrates the traditional plug-and-play method into the diffusion sampling framework. Compared to plug-and-play IR methods that rely on discriminative Gaussian denoisers, DiffPIR is expected to inherit the generative ability of diffusion models. Experimental results on three representative IR tasks, including super-resolution, image deblurring, and inpainting, demonstrate that DiffPIR achieves state-of-the-art performance on both the FFHQ and ImageNet datasets in terms of reconstruction faithfulness and perceptual quality with no more than 100 NFEs.

Setting Up

Clone and Install

git clone https://github.com/yuanzhi-zhu/DiffPIR.git
cd DiffPIR
pip install -r requirements.txt

for motion blur, please download https://github.com/LeviBorodenko/motionblur to the DiffPIR folder.

Model Download

links to model checkpoints can be found in ./model_zoo/README.md

you can also download with:

bash download.sh

Do not forget to rename "ffhq_10m" to "diffusion_ffhq_m" for code consistency.

Inference Code

python main_ddpir_sisr.py # SR
python main_ddpir_deblur.py # deblur
python main_ddpir_inpainting.py # inpainting

Alternatively:

python main_ddpir.py --opt configs/sisr.yaml # SR
python main_ddpir.py --opt configs/deblur.yaml # deblur
python main_ddpir.py --opt configs/inpaint.yaml # inpainting

Train Your Own Diffusion Models

To train a new diffusion model, please follow OpenAI Guided Diffusion.

Brief Introduction

Upon comparison with several earlier iterative image restoration methods, such as USRNet, we found that the diffusion sampling framework offers a more systematic approach to solve data sub-problems and prior sub-problems in an iterative plug and play manner.

Start with the following optimization problem:

$$ \hat{\mathbf{x}} = \mathop{\arg\min}_\mathbf{x} \frac{1}{2\sigma_n^2}|\mathbf{y} - \mathcal{H}(\mathbf{x})|^2 + \lambda \mathcal{P}(\mathbf{x}) $$

In previous works, we can solve this iteratively with:

\begin{align}
\mathbf{{z}}_{k} &= \mathop{\arg\min}_{\mathbf{z}} \frac{1}{2(\sqrt{\lambda/\mu})^2}\|\mathbf{z}-\mathbf{x}_{k}\|^2  + \mathcal{P}(\mathbf{z}) \\
\mathbf{{x}}_{k-1} &= \mathop{\arg\min}_{\mathbf{x}}  \|\mathbf{y} - \mathcal{H}(\mathbf{x})\|^2 + \mu\sigma_n^2\|\mathbf{x}-\mathbf{{z}}_{k} \|^2 
\end{align}

In our study, we address the optimization problem stated above with an altered schedule instead:

\begin{align}
\mathbf{{x}}_{0}^{(t)}&=\mathop{\arg\min}_{\mathbf{z}} {\frac{1}{2\bar{\sigma}_t^2}\|\mathbf{z}-\mathbf{x}_{t}\|^2}  + {\mathcal{P}(\mathbf{z})}\\
\mathbf{\hat{x}}_{0}^{(t)}&=\mathop{\arg\min}_{\mathbf{x}}  \|\mathbf{y} - \mathcal{H}(\mathbf{x})\|^2 + \rho_t\|\mathbf{x}-\mathbf{{x}}_{0}^{(t)} \|^2 \\
\mathbf{x}_{t-1} &\longleftarrow \mathbf{\hat{x}}_{0}^{(t)}
\end{align}

where the data sub-problem is treated as a denoising issue where diffusion models are utilized as generative denoisers.

Figure 1. Illustration of our plug-and-play sampling method.

To detail further, as demonstrated in Figure 1, at every timestep $t$, we initially estimate $\mathbf{x}^{(t)}_0$ from $\mathbf{x}_{t}$ through denoising, using the off-the-shelf unconditional pre-trained diffusion models. Following this, we address the data sub-problem to acquire an updated $\mathbf{\hat x}^{(t)}_0$ (indicated by the red line).

The overall plug and play sampling algorithm can be summarized as follows:

For a more comprehensive understanding, feel free to check out the related paper or slides.

Results

Quantitative Results

Qualitative Results

Citation

If you find this repo helpful, please cite:

@inproceedings{zhu2023denoising, % DiffPIR
      title={Denoising Diffusion Models for Plug-and-Play Image Restoration},
      author={Yuanzhi Zhu and Kai Zhang and Jingyun Liang and Jiezhang Cao and Bihan Wen and Radu Timofte and Luc Van Gool},
      booktitle={IEEE Conference on Computer Vision and Pattern Recognition Workshops (NTIRE)},
      year={2023},
}

Acknowledgments

This work was partly supported by the ETH Zurich General Fund (OK), the Alexander von Humboldt Foundation and the Huawei Fund.