📄 Read our paper on arXiv

This project introduces the DiffusionReward method, which enhances blind face restoration through reward feedback learning. Using DiffBIR+ReFL as an example, we provide inference and training code.

Face Reward Model

DiffusionReward

Before and after comparison

1. Create conda environment

   conda create -n DiffusionReward python=3.10
   conda activate DiffusionReward

2. Install dependencies

   pip install -r requirements.txt

Test Dataset Download

• CelebA-Test, LFW-Test, and Webphoto-Test datasets are from the VQFR public repository
• Download link: VQFR Repository

Download the ControlNet weights for Diffbir+ReFL here.
Place the weight file in the ./weights folder.

Run the following command to perform model validation:

CUDA_VISIBLE_DEVICES=0 python inference.py \
    --task face \
    --diffbir_refl_path ./weights/diffbir_refl.pt \
    --version v1_refl \
    --captioner none \
    --pos_prompt '' \
    --neg_prompt "low quality,blurry,low-resolution,noisy,unsharp,weird textures" \
    --cfg_scale 1.0 \
    --sampler ddim \
    --input ./dataset/celebA/celeba_512_validation_lq/ \
    --output ./output/celeba_diffbir/ \
    --device cuda \
    --precision fp32

• --task: Task type, set to face
• --diffbir_refl_path: Path to DiffBIR+ReFL model weights
• --version: Model version
• --input: Input image path
• --output: Output result save path
• --sampler: Sampler type

1. FFHQ Dataset Preparation

   • First download the FFHQ dataset: FFHQ Dataset
   • Download our prepared text description files: Text Description Files
   • Extract the text description files to the FFHQ image folder
   • Download the images_paths.txt file, which contains relative path information for images

2. Dataset File Structure

   dataset/
   └── FFHQ/
       ├── FFHQ_&_caption/
       │   ├── 00000.png
       │   ├── 00000.txt
       │   ├── ...
       │   ├── 69999.png
       │   └── 69999.txt
       └── image_paths.txt
   

To run the training phase, download the following pre-trained model weights and place them in the ./weights folder:

• Stable Diffusion: Download the pre-trained weights for the SD-2.1 model here.

• DiffBIR: Download the pre-trained weights for the DiffBIR-v1 face model here.

• SwinIR: Download the pre-trained weights for the SwinIR model here.

• FaceReward: Download the pre-trained weights for the FaceReward model here.

Use the following command to start training:

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

Note: Please modify the weight paths in the .yaml file according to your actual file locations.

This work is built upon the excellent foundation provided by the DiffBIR repository. We sincerely thank the authors for their outstanding contribution to the field of blind image restoration.

If you take use of our code or feel our paper is useful for you, please cite our papers:

@misc{wu2025diffusionrewardenhancingblindface,
      title={DiffusionReward: Enhancing Blind Face Restoration through Reward Feedback Learning}, 
      author={Bin Wu and Wei Wang and Yahui Liu and Zixiang Li and Yao Zhao},
      year={2025},
      eprint={2505.17910},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2505.17910}, 
}