This is a Windows tutorial with NVIDIA GPU for running DiffBIR in your command line, but made convenenient ✨
| Total file size | 20GB |
|---|---|
| VRAM | >8GB |
You can watch the YouTube tutorial here
Install Pinokio, we wrote a pinokio file where you just need 1 click to install all of the dependencies.
Then open up Pinokio, go to the top right button "Discover"
Copy the link of this repository, paste the link at the side that says "enter git URL"
And press download
You can find "DiffBIR.git" in your Pinokio list (if you didn't change the save name).
Press Install to download all the dependencies in a conda env
Press the "start" button on Pinokio
enter the following command in the CMD
python gradio_diffbir.py --ckpt weights/general_full_v1.ckpt --config configs/model/cldm.yaml --reload_swinir --swinir_ckpt weights/general_swinir_v1.ckpt
It'll launch a Gradio webUI and you can access it with http://localhost:7860 after it starts running.
The following are a copy of the original repository for references.
This tutorial is made by the team at bycloud AI, feel free to support our endevour in making installation tutorials here!
Xinqi Lin1,*, Jingwen He2,*, Ziyan Chen2, Zhaoyang Lyu2, Ben Fei2, Bo Dai2, Wanli Ouyang2, Yu Qiao2, Chao Dong1,2
1Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences
2Shanghai AI Laboratory
⭐If DiffBIR is helpful for you, please help star this repo. Thanks!🤗
- Visual Results On Real-world Images
- Installation
- Pretrained Models
- Quick Start (gradio demo)
- Inference
- Train
- Update
- TODO
- Python >= 3.9
- CUDA >= 11.3
- PyTorch >= 1.12.1
- xformers == 0.0.16
# clone this repo
git clone https://github.com/XPixelGroup/DiffBIR.git
cd DiffBIR
# create a conda environment with python >= 3.9
conda create -n diffbir python=3.9
conda activate diffbir
conda install pytorch==1.12.1 torchvision==0.13.1 cudatoolkit=11.3 -c pytorch
conda install xformers==0.0.16 -c xformers
# other dependencies
pip install -r requirements.txt| Model Name | Description |
|---|---|
| general_swinir_v1.ckpt | Stage1 model (SwinIR) for general image restoration. |
| general_full_v1.ckpt | Full model for general image restoration. "Full" means it contains both the stage1 and stage2 model. |
| face_swinir_v1.ckpt | Stage1 model (SwinIR) for face restoration. |
| face_full_v1.ckpt | Full model for face restoration. |
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
Download general_full_v1.ckpt and general_swinir_v1.ckpt to weights/ and run the following command.
python inference.py \
--input inputs/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 \
--image_size 512 \
--color_fix_type wavelet --resize_back \
--output results/generalIf you are confused about where the reload_swinir option came from, please refer to the degradation details.
Download face_full_v1.ckpt to weights/ and run the following command.
# for aligned face inputs
python inference_face.py \
--config configs/model/cldm.yaml \
--ckpt weights/face_full_v1.ckpt \
--input inputs/face/aligned \
--steps 50 \
--sr_scale 1 \
--image_size 512 \
--color_fix_type wavelet \
--output results/face/aligned --resize_back \
--has_aligned
# for unaligned face inputs
python inference_face.py \
--config configs/model/cldm.yaml \
--ckpt weights/face_full_v1.ckpt \
--input inputs/face/whole_img \
--steps 50 \
--sr_scale 1 \
--image_size 512 \
--color_fix_type wavelet \
--output results/face/whole_img --resize_backDownload 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]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 --image_size 512 \
--input [img_dir_path] \
--color_fix_type wavelet --resize_back \
--output [output_dir_path] \
--disable_preprocess_modelFor 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.
-
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_folderand split them into training set and validation set automatically. You will get two file lists insave_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.
-
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.
-
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]
- 2023.08.30: Repo is released.
- 2023.09.06: Update colab demo. Thanks to camenduru!:hugs:
- 2023.09.08: Add support for restoring unaligned faces.
- Release code and pretrained models:computer:.
- Update links to paper and project page:link:.
- Release real47 testset:minidisc:.
- Reduce the memory usage of DiffBIR:smiley_cat:.
- Provide HuggingFace demo:notebook:.
- Upload inference code of latent image guidance:page_facing_up:.
- Improve the performance:superhero:.
- Add a patch-based sampling schedule:mag:.
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},
}
This project is released under the Apache 2.0 license.
This project is based on ControlNet and BasicSR. Thanks for their awesome work.
If you have any questions, please feel free to contact with me at linxinqi@tju.edu.cn.