[NeurIPS 2024] DoSSR: Taming Diffusion Prior for Image Super-Resolution with Domain Shift SDEs

📅 Update

2024.11.26: Our code is now open source🎉 and available .

🔎 Visual Results on Real-World Super-Resolution

💻 Installation

git clone https://github.com/AMD-AIG-AIMA/DoSSR
cd DoSSR

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

🌐 Inference

The model is available in two versions, supporting multi-step generation and single-step generation, respectively. The pre-trained model weights are available at google drive. The difference between the multi-step and single-step generation models is the preprocessing module: Real-ESRNet for multi-step and Real-ESRGAN for single-step.

Multi-step Generation

python inference.py \
--input [path/to/input_images] \
--config configs/model/cldm_v21.yaml \
--ckpt [path/to/model_weights(dossr_default.ckpt)] \
--steps 5 \  
--sr_scale 4 \
--color_fix_type wavelet \
--output [path/to/output_folder] \
--device cuda 

One-step Generation

python inference.py \
--input [path/to/input_images] \
--config configs/model/cldm_v21.yaml \
--ckpt [path/to/model_weights(dossr_onestep.ckpt)] \
--steps 1 \
--sr_scale 4 \
--color_fix_type wavelet \
--output [path/to/output_folder] \
--device cuda 

⌛ Train Model

Step 1: Datasets Preparation

Generate the file path list for the training and validation sets.

python scripts/make_file_list.py \
--img_folder  [path/to/datasets] \
--val_size [val_size] \
--save_folder [path/to/save_folder] \
--follow_links 

Step 2: Create the initial model weights

1. Download the pretrained Stable Diffusion models from HuggingFace and the preprocessing model weights from Real-ESRGAN. We use Real-ESRNet(RealESRNet_x4plus.pth) as the preprocessing model during training, as the artifacts introduced by the preprocessing model trained with GAN loss may incorrectly guide the generation model.
2. Use the following script to merge these two weights.

python scripts/make_init_weight.py \
--cldm_config configs/model/cldm_v21.yaml \
--sd_weight [path/to/sd_weights] \
--preprocess_model_weight [path/to/preprocess_model_weights] \
--output [path/to/save_weights]

[Optinal Step]: Pre-crop training samples to accelerate I/O.

python scripts/sample_dataset.py \
--config configs/dataset/general_deg_realesrgan_train.yaml \
--sample_size [sample_size] \
--show_gt \
--output [path/to/save_folder]

Step 3: Start Training

python scripts/train.py --config configs/train_cldm.yaml

🥰 Acknowledgement

Our code is based on DiffBIR .Thanks for their awesome work.

📧 Contact

If you have any questions, please feel free to contact: qinpeng.cui@amd.com, yixuanl@amd.com

📝 Citation

If you find this method and/or code useful, please consider citing

@inproceedings{cuitaming,
  title={Taming Diffusion Prior for Image Super-Resolution with Domain Shift SDEs},
  author={Cui, Qinpeng and Liu, Yixuan and Zhang, Xinyi and Bao, Qiqi and Liao, Qingmin and Tian, Lu and Liu, Zicheng and Wang, Zhongdao and Barsoum, Emad and others},
  booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems}
}