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

-## Visual Results
+## 👀 Visual Results
### Comparisons on Real-World SR

@@ -30,6 +36,144 @@ Update

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.
+
+
+## ⚙ 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