DiffBIR Documentation

Use inference.py to perform super-resolution. You can use either the ECCV paper version (v2) or the updated v2.1 version. v2.1 supports the llava captioner for better results.

# 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

To run DiffBIR on MacOS using either CPU or MPS (Metal Performance Shaders) acceleration, follow these steps:

1. Install Torch (Preview/Nightly version): MPS acceleration requires MacOS 12.3+.
2. Configure requirements: Do not include triton or xformers in your installation, as they are CUDA-specific.
3. Run Inference: Use the --device cpu or --device mps flag when running the inference script. Using mps can accelerate inference.

To use a custom trained model, set --version custom and provide the training configuration path via --train_cfg and the checkpoint path via --ckpt.

python -u inference.py \
--upscale 4 \
--version custom \
--train_cfg [path/to/training/config] \
--ckpt [path/to/saved/checkpoint] \
--captioner llava \
--cfg_scale 8 \
--noise_aug 0 \
--input inputs/demo/bsr \
--output results/custom_demo_bsr

If you have limited VRAM, you can enable tiled sampling by adding specific flags to your inference.py command. This allows for super-resolution with large scale factors on lower-end hardware.

# 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

To install DiffBIR, clone the repository and set up a Conda environment with Python 3.10. The project is optimized for PyTorch 2.2.2 to utilize built-in memory-efficient attention.

Note for older GPUs: If your GPU is incompatible with the latest PyTorch, you can downgrade to pytorch 1.13.1+cu116 and install xformers 0.0.16 as an alternative.

# 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

Windows users may encounter issues installing the triton package.

To avoid these issues, you can choose to run DiffBIR on CPU by omitting the installation of xformers and triton.

If you require CUDA support, you must resolve triton installation issues (refer to the project's issue tracker for community solutions).

You can interact with DiffBIR through a Gradio-based web interface. For users with limited VRAM, you can set the captioner option to ram or none to reduce memory usage.

# For low-VRAM users, set captioner to ram or none
python run_gradio.py --captioner llava

Stage 2 provides generative capabilities using Stable Diffusion.

1. Download Weights: Download Stable Diffusion v2.1 (e.g., v2-1_512-ema-pruned.ckpt).
2. Generate file list: Create a training file list (validation is not currently supported in Stage 2).
3. Configure: Fill in configs/train/train_stage2.yaml.
4. Launch: Use accelerate launch to start training.

Stage 1 involves training a SwinIR model for degradation removal.

1. Generate file lists: Create text files containing paths to your training and validation images.

   find [img_dir] -type f > files.list
   shuf files.list > files_shuf.list
   head -n [train_size] files_shuf.list > files_shuf_train.list
   tail -n +[train_size + 1] files_shuf.list > files_shuf_val.list

2. Configure: Fill in configs/train/train_stage1.yaml.
3. Launch: Use accelerate launch to start training.

accelerate launch train_stage1.py --config configs/train/train_stage1.yaml

Use the face_background task for unaligned face restoration (restoring the whole image including the background).

# 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

Use the face task for aligned face restoration. Note that upscale is typically set to 1 for this task.

# 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

Use the denoise task to remove noise from images.

# 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