github.com/XPixelGroup/DiffBIR @v2.1.0 sqlite

<img src="https://github.com/XPixelGroup/DiffBIR/raw/v2.1.0/assets/logo.png" width="400">

DiffBIR: Towards Blind Image Restoration with Generative Diffusion Prior

Paper | Project Page

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

¹Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences

²Shanghai AI Laboratory

³The Chinese University of Hong Kong

<img src="https://github.com/XPixelGroup/DiffBIR/raw/v2.1.0/assets/teaser.png">

---

<img src="https://github.com/XPixelGroup/DiffBIR/raw/v2.1.0/assets/pipeline.png">

:star:If DiffBIR is helpful for you, please help star this repo. Thanks!:hugs:

:book:Table Of Contents

• Update
• Visual Results On Real-world Images
• TODO
• Installation
• Quick Start
• Pretrained Models
• Inference
• Train

:new:Update

• 2024.11.27: ✅ Release DiffBIR v2.1, including a new model trained on unsplash dataset with LLaVA-generated captions, more samplers, better tiled-sampling support and so on. Check release note for details.
• 2024.04.08: ✅ Release everything about our updated manuscript, including (1) a new model trained on subset of laion2b-en and (2) a more readable code base, etc. DiffBIR is now a general restoration pipeline that could handle different blind image restoration tasks with a unified generation module.
• 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! :rocket: Try it!
• 2023.09.13: :rocket: 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: :rocket: Update colab demo. Thanks to camenduru!:hugs:
• 2023.08.30: This repo is released.

:eyes:Visual Results On Real-world Images

Blind Image Super-Resolution

Blind Face Restoration

:star: Face and the background enhanced by DiffBIR.

Blind Image Denoising

8x Blind Super-Resolution With Tiled Sampling

I often think of Bag End. I miss my books and my arm chair, and my garden. See, that's where I belong. That's home. --- Bilbo Baggins

:climbing:TODO

• [x] Release code and pretrained models :computer:.
• [x] Update links to paper and project page :link:.
• [x] Release real47 testset :minidisc:.
• [ ] Provide webui.
• [x] Reduce the vram usage of DiffBIR :fire::fire::fire:.
• [ ] Provide HuggingFace demo :notebook:.
• [x] Add a patch-based sampling schedule :mag:.
• [x] Upload inference code of latent image guidance :page_facing_up:.
• [x] Improve the performance :superhero:.
• [x] Support MPS acceleration for MacOS users.
• [ ] DiffBIR-turbo :fire::fire::fire:.
• [x] Speed up inference, such as using fp16/bf16, torch.compile :fire::fire::fire:.

:gear:Installation

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

# create environment
conda create -n diffbir python=3.10
conda activate diffbir
pip install -r requirements.txt

Our new code is based on pytorch 2.2.2 for the built-in support of memory-efficient attention. If you are working on a GPU that is not compatible with the latest pytorch, just downgrade pytorch to 1.13.1+cu116 and install xformers 0.0.16 as an alternative.

:flight_departure:Quick Start

Run the following command to interact with the gradio website.

python run_gradio.py

<kbd><img src="https://github.com/XPixelGroup/DiffBIR/raw/v2.1.0/assets/gradio.png"></img></kbd>

:dna:Pretrained Models

Here we list pretrained weight of stage 2 model (IRControlNet) and our trained SwinIR, which was used for degradation removal during the training of stage 2 model.

Model Name	Description	HuggingFace	BaiduNetdisk	OpenXLab
v2.1.pt	IRControlNet trained on filtered unsplash	download	N/A	N/A
v2.pth	IRControlNet trained on filtered laion2b-en	download	download

(pwd: xiu3) | download | | v1_general.pth | IRControlNet trained on ImageNet-1k | download | download

(pwd: 79n9) | download | | v1_face.pth | IRControlNet trained on FFHQ | download | download

(pwd: n7dx) | download | | codeformer_swinir.ckpt | SwinIR trained on ImageNet-1k | download | download

(pwd: vfif) | download |

During inference, we use off-the-shelf models from other papers as the stage 1 model: BSRNet for BSR, SwinIR-Face used in DifFace for BFR, and SCUNet-PSNR for BID, while the trained IRControlNet remains unchanged for all tasks. Please check code for more details. Thanks for their work!

:crossed_swords:Inference

We provide some examples for inference, check inference.py for more arguments. Pretrained weights will be automatically downloaded. For users with limited VRAM, please run the following scripts with tiled sampling.

Blind Image Super-Resolution

# DiffBIR v2 (ECCV paper version)
python -u inference.py \
--task sr \
--upscale 4 \
--version v2 \
--sampler spaced \
--steps 50 \
--captioner none \
--pos_prompt '' \
--neg_prompt 'low quality, blurry, low-resolution, noisy, unsharp, weird textures' \
--cfg_scale 4 \
--input inputs/demo/bsr \
--output results/v2_demo_bsr \
--device cuda --precision fp32

# DiffBIR v2.1
python -u inference.py \
--task sr \
--upscale 4 \
--version v2.1 \
--captioner llava \
--cfg_scale 8 \
--noise_aug 0 \
--input inputs/demo/bsr \
--output results/v2.1_demo_bsr

Blind Aligned-Face Restoration

# DiffBIR v2 (ECCV paper version)
python -u inference.py \
--task face \
--upscale 1 \
--version v2 \
--sampler spaced \
--steps 50 \
--captioner none \
--pos_prompt '' \
--neg_prompt 'low quality, blurry, low-resolution, noisy, unsharp, weird textures' \
--cfg_scale 4.0 \
--input inputs/demo/bfr/aligned \
--output results/v2_demo_bfr_aligned \
--device cuda --precision fp32

# DiffBIR v2.1
python -u inference.py \
--task face \
--upscale 1 \
--version v2.1 \
--captioner llava \
--cfg_scale 8 \
--noise_aug 0 \
--input inputs/demo/bfr/aligned \
--output results/v2.1_demo_bfr_aligned

Blind Unaligned-Face Restoration

# DiffBIR v2 (ECCV paper version)
python -u inference.py \
--task face_background \
--upscale 2 \
--version v2 \
--sampler spaced \
--steps 50 \
--captioner none \
--pos_prompt '' \
--neg_prompt 'low quality, blurry, low-resolution, noisy, unsharp, weird textures' \
--cfg_scale 4.0 \
--input inputs/demo/bfr/whole_img \
--output results/v2_demo_bfr_unaligned \
--device cuda --precision fp32

# DiffBIR v2.1
python -u inference.py \
--task face_background \
--upscale 2 \
--version v2.1 \
--captioner llava \
--cfg_scale 8 \
--noise_aug 0 \
--input inputs/demo/bfr/whole_img \
--output results/v2.1_demo_bfr_unaligned

Blind Image Denoising

# DiffBIR v2 (ECCV paper version)
python -u inference.py \
--task denoise \
--upscale 1 \
--version v2 \
--sampler spaced \
--steps 50 \
--captioner none \
--pos_prompt '' \
--neg_prompt 'low quality, blurry, low-resolution, noisy, unsharp, weird textures' \
--cfg_scale 4.0 \
--input inputs/demo/bid \
--output results/v2_demo_bid \
--device cuda --precision fp32

# DiffBIR v2.1
python -u inference.py \
--task denoise \
--upscale 1 \
--version v2.1 \
--captioner llava \
--cfg_scale 8 \
--noise_aug 0 \
--input inputs/demo/bid \
--output results/v2.1_demo_bid

Other options

Tiled sampling

Add the following arguments to enable tiled sampling:

[command...] \
# tiled inference for stage-1 model
--cleaner_tiled \
--cleaner_tile_size 256 \
--cleaner_tile_stride 128 \
# tiled inference for VAE encoding
--vae_encoder_tiled \
--vae_encoder_tile_size 256 \
# tiled inference for VAE decoding
--vae_decoder_tiled \
--vae_decoder_tile_size 256 \
# tiled inference for diffusion process
--cldm_tiled \
--cldm_tile_size 512 \
--cldm_tile_stride 256

Tiled sampling supports super-resolution with a large scale factor on low-VRAM graphics cards. Our tiled sampling is built upon mixture-of-diffusers and Tiled-VAE. Thanks for their work! <!--

Restoration Guidance

Restoration 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 -u inference.py \
--version v2 \
--task sr \
--upscale 4 \
--cfg_scale 4.0 \
--input inputs/demo/bsr \
--guidance --g_loss w_mse --g_scale 0.5 --g_space rgb \
--output results/demo_bsr_wg \
--device cuda

Y