GitHub - wyf0912/SinSR: [CVPR 2024] SinSR: Diffusion-Based Image Super-Resolution in a Single Step (original) (raw)

[CVPR 24］SinSR: Diffusion-Based Image Super-Resolution in a Single Step

Welcome! This is the official implementation of the paper "SinSR: Diffusion-Based Image Super-Resolution in a Single Step".

• Yufei Wang, Wenhan Yang, Xinyuan Chen, Yaohui Wang, Lanqing Guo, Lap-Pui Chau, Ziwei Liu, Yu Qiao, Alex C. Kot, Bihan Wen
• 1^11 Nanyang Technological University, 2^22 Peng Cheng Laboratory, 3^33 Shanghai Artificial Intelligence Laboratory, 4^44 The Hong Kong Polytechnic University

🐢 Requirements

• Python 3.10, Pytorch 2.1.2, xformers 0.0.23
• More detail (See environment.yml) A suitable conda environment named resshift can be created and activated with:

conda env create -n SinSR python=3.10
conda activate SinSR
pip install -r requirements.txt

or

conda env create -f environment.yml
conda activate SinSR

🐳 Demo

You can try our method through an online demo:

(The time taken for the initial run of the model includes loading the model. Besides, it includes a significant amount of time overhead apart from the algorithms itself, e.g., I/O cost, and web frameworks.)

🚀 Fast Testing

python3 inference.py -i [image folder/image path] -o [result folder] --ckpt weights/SinSR_v1.pth --scale 4 --one_step

Run it on Colab

You can run the code on Google Colab by clicking on the following link:

Requirements

🐬 Reproducing the results in the paper

Results in Table 1

• Real data for image super-resolution: RealSet65 | RealSR
• Test the model

Results on RealSet65

python inference.py -i testdata/RealSet65 -o results/SinSR/RealSet65 --scale 4 --ckpt weights/SinSR_v1.pth --one_step ## Re-evaulated on a RTX3090 # clipiqa: 0.72046 # musiq: 62.25337

Results on RealSR

python inference.py -i testdata/RealSet65 -o results/SinSR/RealSR --scale 4 --ckpt weights/SinSR_v2.pth --one_step ## Re-evaulated on a RTX3090 ### Similar to ResShift, this model is obtained by early stop # clipiqa: 0.69152 # musiq: 61.43469

If you are running on a GPU with limited memory, you could reduce the patch size by setting --chop_size 256 to avoid out of memory. However, this will slightly degrade the performance.

Results on RealSet65

python inference.py -i testdata/RealSet65 -o results/SinSR/RealSet65 --scale 4 --ckpt weights/SinSR_v1.pth --one_step --chop_size 256 --task SinSR

Results on RealSR

python inference.py -i testdata/RealSR -o results/SinSR/RealSR --scale 4 --ckpt weights/SinSR_v2.pth --one_step --chop_size 256 --task SinSR

Results in Table 2

• Download the image ImageNet-Test (Link) to the testdata folder.
• Unzip the downloaded dataset.
• Test the model

python inference.py -i testdata/imagenet256/lq/ -o results/SinSR/imagenet -r testdata/imagenet256/gt/ --scale 4 --ckpt weights/SinSR_v1.pth --one_step ## Re-evaulated on a RTX3090 # clipiqa: 0.60969 # musiq: 53.51805 # psnr: 24.70071 # lpips: 0.21882 # ssim: 0.66364

✈️ Training

Preparing stage

1. Download the necessary pre-trained model, i.e., pretrained ResShift, and Autoencoder. This can be achieved by inferece using ResShift and the needed models will be downloaded automatically.

Method 1

python3 app.py # Select the model to ResShift in the webpage

Method 2

python inference.py --task realsrx4 -i [image folder/image path] -o [result folder] --scale 4 # Inference using ResShift

1. Adjust the data path in the config file. Specifically, correct and complete paths in files of traindata
2. Adjust batchsize according your GPUS.
   • configs.train.batch: [training batchsize, validation btatchsize]
   • configs.train.microbatch: total batchsize = microbatch * #GPUS * num_grad_accumulation

Train the model

python3 main_distill.py --cfg_path configs/SinSR.yaml --save_dir logs/SinSR

We find that the model can converge very quickly, e.g., a few thousand iterations. Therefore, we believe that the proposed method could be applied to other diffuson-based SR models and encourage a try if you are interested.

❤️ Acknowledgement

This project is based on ResShift. Thanks for the help from the author.

⭐ Citation

Please cite our paper if you find our work useful. Thanks!

@inproceedings{wang2024sinsr,
  title={SinSR: diffusion-based image super-resolution in a single step},
  author={Wang, Yufei and Yang, Wenhan and Chen, Xinyuan and Wang, Yaohui and Guo, Lanqing and Chau, Lap-Pui and Liu, Ziwei and Qiao, Yu and Kot, Alex C and Wen, Bihan},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={25796--25805},
  year={2024}
}

📧 Contact

If you have any questions, please feel free to contact me via yufei001@ntu.edu.sg.