[Paper] [Project Page]

This repository contains PyTorch implementation for Generative Diffusion Prior for Unified Image Restoration and Enhancement.

This is the repo is based on the open-source repo for Guided Diffusion. Run the following command the install the guided-diffusion package:

pip install -e .

Download Checkpoints and Data

We use a single single unconditional DDPM pre-trained on ImageNet to produce diverse and high-fidelity outputs for unified image restoration and enhancement. Download pretrained DDPMs on ImageNet-256(uncond) from this page. Then download pre-generated ImageNet-64 (BigGAN-deep), LSUN-Bedroom (StyleGAN) and LSUN-Cat (StyleGAN2) images from this page this page.

We use 1,000 images from the ImageNet validation set for comparison with other methods. The list of images is taken from https://github.com/XingangPan/deep-generative-prior/

Generative Diffusion Prior (GDP-x0) for linear problems

cd scripts

• Super-resolution:

MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
python sample_x0_sr4.py  $MODEL_FLAGS --use_img_for_guidance --start_from_scratch --save_png_files

• Deblurring:

MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
python sample_x0_blur.py  $MODEL_FLAGS --use_img_for_guidance --start_from_scratch --save_png_files

• Impainting:

MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
python sample_x0_linear.py  $MODEL_FLAGS --deg inp25 --use_img_for_guidance --start_from_scratch --save_png_files

• Colorization:

MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
python sample_x0_color.py  $MODEL_FLAGS --use_img_for_guidance --start_from_scratch --save_png_files

Generative Diffusion Prior (GDP) for non-linear, blind, multi-guidance problems

• Low light enhancement:

MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
python sample_x0_enhancement_low_light.py  $MODEL_FLAGS --use_img_for_guidance --start_from_scratch --save_png_files

• HDR recovery:

MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
python sample_x0_enhancement_HDR.py  $MODEL_FLAGS --use_img_for_guidance --start_from_scratch --save_png_files

Generative Diffusion Prior (GDP-xt) for linear problems

cd scripts

• Super-resolution:

MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
python sample_sr4.py  $MODEL_FLAGS --use_img_for_guidance --start_from_scratch --save_png_files

• Deblurring:

MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
python sample_blur.py  $MODEL_FLAGS --use_img_for_guidance --start_from_scratch --save_png_files

• Impainting:

MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
python sample_inp.py  $MODEL_FLAGS --deg inp25 --use_img_for_guidance --start_from_scratch --save_png_files

• Colorization:

MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
python sample_color.py  $MODEL_FLAGS --use_img_for_guidance --start_from_scratch --save_png_files

Generative Diffusion Prior (GDP-x0-v1) for linear problems

cd scripts

• Super-resolution:

MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
python sample_x0_sr4_v1.py  $MODEL_FLAGS --use_img_for_guidance --start_from_scratch --save_png_files

• Deblurring:

MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
python sample_x0_blur_v1.py  $MODEL_FLAGS --use_img_for_guidance --start_from_scratch --save_png_files

• Impainting:

MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
python sample_x0_inp_v1.py  $MODEL_FLAGS --deg inp25 --use_img_for_guidance --start_from_scratch --save_png_files

• Colorization:

MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
python sample_x0_color_v1.py  $MODEL_FLAGS --use_img_for_guidance --start_from_scratch --save_png_files

References and Acknowledgements

@inproceedings{fei2022generative,
    title={Generative Diffusion Prior for Unified Image Restoration and Enhancement},
    author={Ben Fei, Zhaoyang Lyu, Liang Pan, Junzhe Zhang, Weidong Yang, Tianyue Luo, Bo Zhang, Bo Dai},
    booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
    year={2023}
}

This implementation is based on / inspired by:

• https://github.com/openai/guided-diffusion (the Guided Diffusion repo),
• https://github.com/bahjat-kawar/ddrm (the DDRM repo), and
• https://github.com/ZhaoyangLyu/Early_Stopped_DDPM (the Early_Stopped_DDPM repo)