DiffBIR: A Python repository from alphapro-pro

DiffBIR: Towards Blind Image Restoration with Generative Diffusion Prior

Paper | Project Page

Xinqi Lin^1,*, Jingwen He^2,*, Ziyan Chen², Zhaoyang Lyu², Ben Fei², Bo Dai², Wanli Ouyang², Yu Qiao², Chao Dong^1,2

¹Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences
²Shanghai AI Laboratory

⭐If DiffBIR is helpful for you, please help star this repo. Thanks!🤗

👀Visual Results On Real-world Images

General Image Restoration

Face Image Restoration

⭐ Face and the background enhanced by DiffBIR.

🆕Update

2023.09.19: ✅ Add support for Apple Silicon! Check installation_xOS.md to work with CPU/CUDA/MPS device!
2023.09.14: ✅ Integrate a patch-based sampling strategy (mixture-of-diffusers). Try it! Here is an example with a resolution of 2396 x 1596. GPU memory usage will continue to be optimized in the future and we are looking forward to your pull requests!
2023.09.14: ✅ Add support for background upsampler(DiffBIR/RealESRGAN) in face enhancement! 🚀 Try it!
2023.09.13: 🚀 Provide online demo (DiffBIR-official) in OpenXLab, which integrates both general model and face model. Please have a try! camenduru also implements an online demo, thanks for his work.:hugs:
2023.09.12: ✅ Upload inference code of latent image guidance and release real47 testset.
2023.09.08: ✅ Add support for restoring unaligned faces.
2023.09.06: 🚀 Update colab demo. Thanks to camenduru!:hugs:
2023.08.30: This repo is released.

🧗TODO

Release code and pretrained models:computer:.
Update links to paper and project page:link:.
Release real47 testset:minidisc:.
Provide webui and reduce the memory usage of DiffBIR:fire::fire::fire:.
Provide HuggingFace demo:notebook::fire::fire::fire:.
Add a patch-based sampling schedule:mag:.
Upload inference code of latent image guidance:page_facing_up:.
Improve the performance:superhero:.
Support MPS acceleration for MacOS users.

⚙️Installation

# clone this repo
git clone https://github.com/XPixelGroup/DiffBIR.git
cd DiffBIR

# create an environment with python >= 3.9
conda create -n diffbir python=3.9
conda activate diffbir
pip install -r requirements.txt

Note the installation is only compatible with Linux users. If you are working on different platforms, please check xOS Installation.

🧬Pretrained Models

Model Name	Description	HuggingFace	BaiduNetdisk	OpenXLab
general_swinir_v1.ckpt	Stage1 model (SwinIR) for general image restoration.	download	download (pwd: v3v6)	download
general_full_v1.ckpt	Full model for general image restoration. "Full" means it contains both the stage1 and stage2 model.	download	download (pwd: 86zi)	download
face_swinir_v1.ckpt	Stage1 model (SwinIR) for face restoration.	download	download (pwd: xk5u)	download
face_full_v1.ckpt	Full model for face restoration.	download	download (pwd: ov8i)	download

🛫Quick Start

Download general_full_v1.ckpt and general_swinir_v1.ckpt to weights/, then run the following command to interact with the gradio website.

python gradio_diffbir.py \
--ckpt weights/general_full_v1.ckpt \
--config configs/model/cldm.yaml \
--reload_swinir \
--swinir_ckpt weights/general_swinir_v1.ckpt \
--device cuda

⚔️Inference

Full Pipeline (Remove Degradations & Refine Details)

General Image

Download general_full_v1.ckpt and general_swinir_v1.ckpt to weights/ and run the following command.

python inference.py \
--input inputs/demo/general \
--config configs/model/cldm.yaml \
--ckpt weights/general_full_v1.ckpt \
--reload_swinir --swinir_ckpt weights/general_swinir_v1.ckpt \
--steps 50 \
--sr_scale 4 \
--color_fix_type wavelet \
--output results/demo/general \
--device cuda [--tiled --tile_size 512 --tile_stride 256]

Remove the brackets to enable tiled sampling. If you are confused about where the reload_swinir option came from, please refer to the degradation details.

Face Image

The face_full_v1.ckpt will be downloaded from HuggingFace automatically.

# for aligned face inputs
python inference_face.py \
--input inputs/demo/face/aligned \
--sr_scale 1 \
--output results/demo/face/aligned \
--has_aligned \
--device cuda

# for unaligned face inputs
python inference_face.py \
--input inputs/demo/face/whole_img \
--sr_scale 2 \
--output results/demo/face/whole_img \
--bg_upsampler DiffBIR \
--device cuda

Latent Image Guidance (Quality-fidelity trade-off)

Latent image guidance is used to achieve a trade-off bwtween quality and fidelity. We default to closing it since we prefer quality rather than fidelity. Here is an example:

python inference.py \
--input inputs/demo/general \
--config configs/model/cldm.yaml \
--ckpt weights/general_full_v1.ckpt \
--reload_swinir --swinir_ckpt weights/general_swinir_v1.ckpt \
--steps 50 \
--sr_scale 4 \
--color_fix_type wavelet \
--output results/demo/general \
--device cuda \
--use_guidance --g_scale 400 --g_t_start 200

You will see that the results become more smooth.

Only Stage1 Model (Remove Degradations)

Download general_swinir_v1.ckpt, face_swinir_v1.ckpt for general, face image respectively, and run the following command.

python scripts/inference_stage1.py \
--config configs/model/swinir.yaml \
--ckpt [swinir_ckpt_path] \
--input [lq_dir] \
--sr_scale 1 --image_size 512 \
--output [output_dir_path]

Only Stage2 Model (Refine Details)

Since the proposed two-stage pipeline is very flexible, you can utilize other awesome models to remove degradations instead of SwinIR and then leverage the Stable Diffusion to refine details.

# step1: Use other models to remove degradations and save results in [img_dir_path].

# step2: Refine details of step1 outputs.
python inference.py \
--config configs/model/cldm.yaml \
--ckpt [full_ckpt_path] \
--steps 50 --sr_scale 1 \
--input [img_dir_path] \
--color_fix_type wavelet \
--output [output_dir_path] \
--disable_preprocess_model \
--device cuda

🌠Train

Degradation Details

For general image restoration, we first train both the stage1 and stage2 model under codeformer degradation to enhance the generative capacity of the stage2 model. In order to improve the ability for degradation removal, we train another stage1 model under Real-ESRGAN degradation and utilize it during inference.

For face image restoration, we adopt the degradation model used in DifFace for training and directly utilize the SwinIR model released by them as our stage1 model.

Data Preparation

Generate file list of training set and validation set.
```
python scripts/make_file_list.py \
--img_folder [hq_dir_path] \
--val_size [validation_set_size] \
--save_folder [save_dir_path] \
--follow_links
```
This script will collect all image files in img_folder and split them into training set and validation set automatically. You will get two file lists in save_folder, each line in a file list contains an absolute path of an image file:
```
save_folder
├── train.list # training file list
└── val.list   # validation file list
```
Configure training set and validation set.

For general image restoration, fill in the following configuration files with appropriate values.
- training set and validation set for CodeFormer degradation.
- training set and validation set for Real-ESRGAN degradation.
For face image restoration, fill in the face training set and validation set configuration files with appropriate values.

Train Stage1 Model

Configure training-related information.

Fill in the configuration file of training with appropriate values.
Start training.
```
python train.py --config [training_config_path]
```
💡:Checkpoints of SwinIR will be used in training stage2 model.

Train Stage2 Model

Download pretrained Stable Diffusion v2.1 to provide generative capabilities.

wget https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-ema-pruned.ckpt --no-check-certificate

Create the initial model weights.

python scripts/make_stage2_init_weight.py \
--cldm_config configs/model/cldm.yaml \
--sd_weight [sd_v2.1_ckpt_path] \
--swinir_weight [swinir_ckpt_path] \
--output [init_weight_output_path]

You will see some outputs which show the weight initialization.

Configure training-related information.

Fill in the configuration file of training with appropriate values.

Start training.

python train.py --config [training_config_path]

Citation

Please cite us if our work is useful for your research.

@article{2023diffbir,
  author    = {Xinqi Lin, Jingwen He, Ziyan Chen, Zhaoyang Lyu, Ben Fei, Bo Dai, Wanli Ouyang, Yu Qiao, Chao Dong},
  title     = {DiffBIR: Towards Blind Image Restoration with Generative Diffusion Prior},
  journal   = {arxiv},
  year      = {2023},
}

License

This project is released under the Apache 2.0 license.

Acknowledgement

This project is based on ControlNet and BasicSR. Thanks for their awesome work.

Contact

If you have any questions, please feel free to contact with me at linxinqi@tju.edu.cn.

alphapro-pro/DiffBIR