diff --git a/README.md b/README.md index afb213e..e90e669 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,11 @@ +

+ +

+ ## Improving the Stability of Diffusion Models for Content Consistent Super-Resolution -Paper, codes and pretained models will be released soon. + [Lingchen Sun](https://scholar.google.com/citations?hl=zh-CN&tzom=-480&user=ZCDjTn8AAAAJ)1,2 @@ -13,16 +17,18 @@ Paper, codes and pretained models will be released soon. 1The Hong Kong Polytechnic University, 2OPPO Research Institute -Update +## ⏰ Update +- **2024.1.4**: Code and the model for real-world SR are released. +- **2024.1.3**: Paper is released. - **2023.12.23**: Repo is released. :star: If CCSR is helpful to your images or projects, please help star this repo. Thanks! :hugs: -## Overview Framework +## 🌟 Overview Framework ![ccsr](figs/framework.png) -## Visual Results +## 👀 Visual Results ### Comparisons on Real-World SR ![ccsr](figs/realworld.png) @@ -30,6 +36,144 @@ Update ![ccsr](figs/bicubic.png) For more comparisons, please refer to our paper for details. +## 📝 Quantitative comparisons +We propose new stability metrics, namely global standard deviation (G-STD) and local standard deviation (L-STD), to measure the image-level and pixel-level variations of the SR results for diffusion-based methods. + +More details about G-STD and L-STD can be found in our paper. + +![ccsr](figs/table.png) +## ⚙ Dependencies and Installation +```shell +## git clone this repository +git clone https://github.com/csslc/CCSR.git +cd CCSR + +# create an environment with python >= 3.9 +conda create -n ccsr python=3.9 +conda activate ccsr +pip install -r requirements.txt +``` +## 🍭 Quick Inference +#### Step 1: Download the pretrained models +- Download the pretrained SD-2.1base models from [HuggingFace](https://huggingface.co/stabilityai/stable-diffusion-2-1-base). +- Download the CCSR models from: + +| Model Name | Description | GoogleDrive | OneDive | +|:---------------------|:---------------------------------------------|:--------------------------------------------------------------------------------------|:--------------------------------------------------------------------------------| +| real-world_ccsr.ckpt | CCSR model for real-world image restoration. | [download](https://drive.google.com/drive/folders/1jM1mxDryPk9CTuFTvYcraP2XIVzbPiw_?usp=drive_link) | [download](https://pan.baidu.com/s/1uvSvJgcoL_Knj0h22-9TvA?pwd=v3v6) | +| bicubic_ccsr.ckpt | CCSR model for bicubic image restoration. | download | download | + + +#### Step 2: Prepare testing data +You can put the testing images in the `preset/test_datasets`. + +#### Step 3: Running testing command +``` +python inference_ccsr.py \ +--input preset/datasets/test_datasets \ +--config CCSR-main/configs/model/ccsr_stage2.yaml \ +--ckpt [ccsr_realworld_ckpt_path] \ +--steps 45 \ +--sr_scale 4 \ +--t_max: 0.6667 \ +--t_min: 0.3333 \ +--output experiments/output \ +--device cuda \ +--repeat_times 1 +``` +You can obtain `N` different SR results by setting `repeat_time` as `N` to test the stability of CCSR. The data folder should be like this: + +``` + experiments/output + ├── sample0 # the first group of SR results + └── sample1 # the second group of SR results + ... + └── sampleN # the N-th group of SR results +``` + +## 📏 Evaluation +1. Calculate the Image Quality Assessment for each restored group. + + Fill in the required information in [cal_iqa.py](cal_iqa/cal_iqa.py) and run, then you can obtain the evaluation results in the folder like this: + ``` + log_path + ├── log_name_npy # save the IQA values of each restored group to the npy files + └── log_name.log # log recode + ``` + +2. Calculate the G-STD value for the diffusion-based SR method. + + Fill in the required information in [iqa_G-STD.py](cal_iqa/iqa_G-STD.py) and run, then you can obtain the mean IQA values of N restored groups and G-STD value. + +3. Calculate the L-STD value for the diffusion-based SR method. + + Fill in the required information in [iqa_L-STD.py](cal_iqa/iqa_L-STD.py) and run, then you can obtain the L-STD value. + + +## 🚋 Train + +#### Step1: Prepare training data + +1. Generate file list of training set and validation set. + + ```shell + python scripts/make_file_list.py \ + --img_folder [hq_dir_path] \ + --val_size [validation_set_size] \ + --save_folder [save_dir_path] \ + --follow_links + ``` + + This script will collect all image files in `img_folder` and split them into training set and validation set automatically. You will get two file lists in `save_folder`, each line in a file list contains an absolute path of an image file: + + ``` + save_folder + ├── train.list # training file list + └── val.list # validation file list + ``` + +2. Configure training set and validation set. + + For general image restoration, fill in the following configuration files with appropriate values. + + - [training set](configs/dataset/general_deg_stablesr_realesrgan_train.yaml) and [validation set](configs/dataset/general_deg_stablesr_realesrgan_val.yaml) for **Real-ESRGAN** degradation. + - [training set](configs/dataset/general_deg_bicubic_train.yaml) and [validation set](configs/dataset/general_deg_bicubic_val.yaml) for **Bicubic** degradation. + +#### Step2: Train Stage1 Model +1. Download pretrained [Stable Diffusion v2.1](https://huggingface.co/stabilityai/stable-diffusion-2-1-base) to provide generative capabilities. + + ```shell + wget https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-ema-pruned.ckpt --no-check-certificate + ``` + +2. Create the initial model weights. + + ```shell + python scripts/make_stage2_init_weight.py \ + --cldm_config configs/model/ccsr_stage1.yaml \ + --sd_weight [sd_v2.1_ckpt_path] \ + --output [init_weight_output_path] + ``` + +3. Configure training-related information. + + Fill in the configuration file of [training of stage1](configs/train_ccsr_stage1.yaml) with appropriate settings. + +4. Start training. + + ```shell + python train.py --config configs/train_ccsr_stage1.yaml + ``` + +#### Step3: Train Stage2 Model +1. Configure training-related information. + + Fill in the configuration file of [training of stage2](configs/train_ccsr_stage2.yaml) with appropriate settings. + +2. Start training. + ```shell + python train.py --config configs/train_ccsr_stage2.yaml + ``` ### License This project is released under the [Apache 2.0 license](LICENSE). @@ -43,8 +187,8 @@ The following are BibTeX references: @article{sun2023ccsr, title={Improving the Stability of Diffusion Models for Content Consistent Super-Resolution}, author={Sun, Lingchen and Wu, Rongyuan and Zhang, Zhengqiang and Yong, Hongwei and Zhang, Lei}, - journal={arXiv preprint arXiv:***}, - year={2023} + journal={arXiv preprint arXiv:2401.00877}, + year={2024} } ``` ### Contact diff --git a/figs/bicubic.png b/figs/bicubic.png index 6ddf772..4db36f7 100644 Binary files a/figs/bicubic.png and b/figs/bicubic.png differ diff --git a/figs/bicubic_1.pdf b/figs/bicubic_1.pdf deleted file mode 100644 index e3bdb72..0000000 Binary files a/figs/bicubic_1.pdf and /dev/null differ diff --git a/figs/logo.png b/figs/logo.png new file mode 100644 index 0000000..79c46f3 Binary files /dev/null and b/figs/logo.png differ diff --git a/figs/realworld.png b/figs/realworld.png index 51be4f6..d6fd24b 100644 Binary files a/figs/realworld.png and b/figs/realworld.png differ diff --git a/figs/table.png b/figs/table.png new file mode 100644 index 0000000..220fd4c Binary files /dev/null and b/figs/table.png differ