Upload 7 files

README.md

@@ -0,0 +1,339 @@
+---
+license: apache-2.0
+tags:
+- Image-to-Image
+---
+![1](https://github.com/open-mmlab/mmdetection/assets/95841578/1298bd73-ac5e-4275-a2f6-0a42f2430d79)
+## DiffBIR: Towards Blind Image Restoration with Generative Diffusion Prior
+
+[Paper](https://arxiv.org/abs/2308.15070) | [Project Page](https://0x3f3f3f3fun.github.io/projects/diffbir/)
+
+![visitors](https://visitor-badge.laobi.icu/badge?page_id=XPixelGroup/DiffBIR) [![Open in OpenXLab](https://cdn-static.openxlab.org.cn/app-center/openxlab_app.svg)](https://openxlab.org.cn/apps/detail/linxinqi/DiffBIR-official) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/camenduru/DiffBIR-colab/blob/main/DiffBIR_colab.ipynb)
+
+[Xinqi Lin](https://0x3f3f3f3fun.github.io/)<sup>1,\*</sup>, [Jingwen He](https://github.com/hejingwenhejingwen)<sup>2,\*</sup>, [Ziyan Chen](https://orcid.org/0000-0001-6277-5635)<sup>2</sup>, [Zhaoyang Lyu](https://scholar.google.com.tw/citations?user=gkXFhbwAAAAJ&hl=en)<sup>2</sup>, [Ben Fei](https://scholar.google.com/citations?user=skQROj8AAAAJ&hl=zh-CN&oi=ao)<sup>2</sup>, [Bo Dai](http://daibo.info/)<sup>2</sup>, [Wanli Ouyang](https://wlouyang.github.io/)<sup>2</sup>, [Yu Qiao](http://mmlab.siat.ac.cn/yuqiao)<sup>2</sup>, [Chao Dong](http://xpixel.group/2010/01/20/chaodong.html)<sup>1,2</sup>
+
+<sup>1</sup>Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences<br><sup>2</sup>Shanghai AI Laboratory
+
+![2](https://github.com/open-mmlab/mmdetection/assets/95841578/63b4899a-4c92-4a08-a2bc-e932320e4166)
+
+:star:If DiffBIR is helpful for you, please help star this repo. Thanks!:hugs:
+
+## Table Of Contents
+
+- [Visual Results On Real-world Images](#visual_results)
+- [Update](#update)
+- [TODO](#todo)
+- [Installation](#installation)
+- [Pretrained Models](#pretrained_models)
+- [Quick Start (gradio demo)](#quick_start)
+- [Inference](#inference)
+- [Train](#train)
+
+## <a name="visual_results"></a>Visual Results On Real-world Images
+
+<!-- <details close>
+<summary>General Image Restoration</summary> -->
+### General Image Restoration
+
+![3](https://github.com/open-mmlab/mmdetection/assets/95841578/0887c64d-ba44-4124-b001-f2576e543226)
+
+<!-- <summary>Face Image Restoration</summary> -->
+### Face Image Restoration
+
+![4](https://github.com/open-mmlab/mmdetection/assets/95841578/8e054d66-050a-4eff-8837-acfa4cffc904)
+
+Face and the background enhanced by DiffBIR.
+
+<!-- </details> -->
+
+## <a name="update"></a>Update
+
+- **2023.09.19**: ✅ Add support for Apple Silicon! Check [installation_xOS.md](assets/docs/installation_xOS.md) to work with **CPU/CUDA/MPS** device!
+- **2023.09.14**: ✅ Integrate a patch-based sampling strategy ([mixture-of-diffusers](https://github.com/albarji/mixture-of-diffusers)). [**Try it!**](#general_image_inference) Here is an [example](https://imgsli.com/MjA2MDA1) with a resolution of 2396 x 1596. GPU memory usage will continue to be optimized in the future and we are looking forward to your pull requests!
+- **2023.09.14**: ✅ Add support for background upsampler(DiffBIR/[RealESRGAN](https://github.com/xinntao/Real-ESRGAN)) in face enhancement! :rocket: [**Try it!**](#unaligned_face_inference)
+- **2023.09.13**: :rocket: Provide online demo (DiffBIR-official) in [OpenXLab](https://openxlab.org.cn/apps/detail/linxinqi/DiffBIR-official), which integrates both general model and face model. Please have a try! [camenduru](https://github.com/camenduru) also implements an online demo, thanks for his work.:hugs:
+- **2023.09.12**: ✅ Upload inference code of latent image guidance and release [real47](inputs/real47) testset.
+- **2023.09.08**: ✅ Add support for restoring unaligned faces.
+- **2023.09.06**: :rocket: Update [colab demo](https://colab.research.google.com/github/camenduru/DiffBIR-colab/blob/main/DiffBIR_colab.ipynb). Thanks to [camenduru](https://github.com/camenduru)!:hugs:
+- **2023.08.30**: This repo is released.
+<!-- - [**History Updates** >]() -->
+
+
+## <a name="todo"></a>TODO
+
+- [x] Release code and pretrained models:computer:.
+- [x] Update links to paper and project page:link:.
+- [x] Release real47 testset:minidisc:.
+- [ ] Provide webui and reduce the memory usage of DiffBIR:fire::fire::fire:.
+- [ ] Provide HuggingFace demo:notebook::fire::fire::fire:.
+- [x] Add a patch-based sampling schedule:mag:.
+- [x] Upload inference code of latent image guidance:page_facing_up:.
+- [ ] Improve the performance:superhero:.
+- [x] Support MPS acceleration for MacOS users.
+
+## <a name="installation"></a>Installation
+<!-- - **Python** >= 3.9
+- **CUDA** >= 11.3
+- **PyTorch** >= 1.12.1
+- **xformers** == 0.0.16 -->
+
+```shell
+# clone this repo
+git clone https://github.com/XPixelGroup/DiffBIR.git
+cd DiffBIR
+
+# create an environment with python >= 3.9
+conda create -n diffbir python=3.9
+conda activate diffbir
+pip install -r requirements.txt
+```
+
+Note the installation is only compatible with **Linux** users. If you are working on different platforms, please check [xOS Installation](assets/docs/installation_xOS.md).
+
+<!-- ```shell
+# clone this repo
+git clone https://github.com/XPixelGroup/DiffBIR.git
+cd DiffBIR
+
+# create a conda environment with python >= 3.9
+conda create -n diffbir python=3.9
+conda activate diffbir
+
+conda install pytorch==1.12.1 torchvision==0.13.1 cudatoolkit=11.3 -c pytorch
+conda install xformers==0.0.16 -c xformers
+
+# other dependencies
+pip install -r requirements.txt
+``` -->
+
+## <a name="pretrained_models"></a>Pretrained Models
+
+| Model Name | Description | HuggingFace | BaiduNetdisk | OpenXLab |
+| :--------- | :---------- | :---------- | :---------- | :---------- |
+| general_swinir_v1.ckpt | Stage1 model (SwinIR) for general image restoration. | [download](https://huggingface.co/lxq007/DiffBIR/resolve/main/general_swinir_v1.ckpt) | [download](https://pan.baidu.com/s/1uvSvJgcoL_Knj0h22-9TvA?pwd=v3v6) (pwd: v3v6) | [download](https://download.openxlab.org.cn/models/linxinqi/DiffBIR/weight//diffbir_general_swinir_v1) |
+| general_full_v1.ckpt | Full model for general image restoration. "Full" means it contains both the stage1 and stage2 model. | [download](https://huggingface.co/lxq007/DiffBIR/resolve/main/general_full_v1.ckpt) | [download](https://pan.baidu.com/s/1gLvW1nvkJStdVAKROqaYaA?pwd=86zi) (pwd: 86zi) | [download](https://download.openxlab.org.cn/models/linxinqi/DiffBIR/weight//diffbir_general_full_v1) |
+| face_swinir_v1.ckpt | Stage1 model (SwinIR) for face restoration. | [download](https://huggingface.co/lxq007/DiffBIR/resolve/main/face_swinir_v1.ckpt) | [download](https://pan.baidu.com/s/1cnBBC8437BJiM3q6suaK8g?pwd=xk5u) (pwd: xk5u) | [download](https://download.openxlab.org.cn/models/linxinqi/DiffBIR/weight//diffbir_face_swinir_v1) |
+| face_full_v1.ckpt | Full model for face restoration. | [download](https://huggingface.co/lxq007/DiffBIR/resolve/main/face_full_v1.ckpt) | [download](https://pan.baidu.com/s/1pc04xvQybkynRfzK5Y8K0Q?pwd=ov8i) (pwd: ov8i) | [download](https://download.openxlab.org.cn/models/linxinqi/DiffBIR/weight//diffbir_face_full_v1) |
+
+## <a name="quick_start"></a>Quick Start
+
+Download [general_full_v1.ckpt](https://huggingface.co/lxq007/DiffBIR/resolve/main/general_full_v1.ckpt) and [general_swinir_v1.ckpt](https://huggingface.co/lxq007/DiffBIR/resolve/main/general_swinir_v1.ckpt) to `weights/`, then run the following command to interact with the gradio website.
+
+```shell
+python gradio_diffbir.py \
+--ckpt weights/general_full_v1.ckpt \
+--config configs/model/cldm.yaml \
+--reload_swinir \
+--swinir_ckpt weights/general_swinir_v1.ckpt \
+--device cuda
+```
+
+<img width="887" alt="5" src="https://github.com/open-mmlab/mmdetection/assets/95841578/36afc84f-61d9-4514-88c8-40eaec557e44">
+
+## <a name="inference"></a>Inference
+
+### Full Pipeline (Remove Degradations & Refine Details)
+
+<a name="general_image_inference"></a>
+#### General Image
+
+Download [general_full_v1.ckpt](https://huggingface.co/lxq007/DiffBIR/resolve/main/general_full_v1.ckpt) and [general_swinir_v1.ckpt](https://huggingface.co/lxq007/DiffBIR/resolve/main/general_swinir_v1.ckpt) to `weights/` and run the following command.
+
+```shell
+python inference.py \
+--input inputs/demo/general \
+--config configs/model/cldm.yaml \
+--ckpt weights/general_full_v1.ckpt \
+--reload_swinir --swinir_ckpt weights/general_swinir_v1.ckpt \
+--steps 50 \
+--sr_scale 4 \
+--color_fix_type wavelet \
+--output results/demo/general \
+--device cuda [--tiled --tile_size 512 --tile_stride 256]
+```
+
+Remove the brackets to enable tiled sampling. If you are confused about where the `reload_swinir` option came from, please refer to the [degradation details](#degradation-details).
+
+#### Face Image
+<!-- Download [face_full_v1.ckpt](https://huggingface.co/lxq007/DiffBIR/resolve/main/face_full_v1.ckpt) to `weights/` and run the following command. -->
+The [face_full_v1.ckpt](https://huggingface.co/lxq007/DiffBIR/resolve/main/face_full_v1.ckpt) will be downloaded from HuggingFace automatically.
+
+```shell
+# for aligned face inputs
+python inference_face.py \
+--input inputs/demo/face/aligned \
+--sr_scale 1 \
+--output results/demo/face/aligned \
+--has_aligned \
+--device cuda
+```
+
+<a name="unaligned_face_inference"></a>
+
+```shell
+# for unaligned face inputs
+python inference_face.py \
+--input inputs/demo/face/whole_img \
+--sr_scale 2 \
+--output results/demo/face/whole_img \
+--bg_upsampler DiffBIR \
+--device cuda
+```
+
+### Latent Image Guidance (Quality-fidelity trade-off)
+
+Latent image guidance is used to achieve a trade-off bwtween quality and fidelity. We default to closing it since we prefer quality rather than fidelity. Here is an example:
+
+```shell
+python inference.py \
+--input inputs/demo/general \
+--config configs/model/cldm.yaml \
+--ckpt weights/general_full_v1.ckpt \
+--reload_swinir --swinir_ckpt weights/general_swinir_v1.ckpt \
+--steps 50 \
+--sr_scale 4 \
+--color_fix_type wavelet \
+--output results/demo/general \
+--device cuda \
+--use_guidance --g_scale 400 --g_t_start 200
+```
+
+You will see that the results become more smooth.
+
+### Only Stage1 Model (Remove Degradations)
+
+Download [general_swinir_v1.ckpt](https://huggingface.co/lxq007/DiffBIR/resolve/main/general_swinir_v1.ckpt), [face_swinir_v1.ckpt](https://huggingface.co/lxq007/DiffBIR/resolve/main/face_swinir_v1.ckpt) for general, face image respectively, and run the following command.
+
+```shell
+python scripts/inference_stage1.py \
+--config configs/model/swinir.yaml \
+--ckpt [swinir_ckpt_path] \
+--input [lq_dir] \
+--sr_scale 1 --image_size 512 \
+--output [output_dir_path]
+```
+
+### Only Stage2 Model (Refine Details)
+
+Since the proposed two-stage pipeline is very flexible, you can utilize other awesome models to remove degradations instead of SwinIR and then leverage the Stable Diffusion to refine details.
+
+```shell
+# step1: Use other models to remove degradations and save results in [img_dir_path].
+
+# step2: Refine details of step1 outputs.
+python inference.py \
+--config configs/model/cldm.yaml \
+--ckpt [full_ckpt_path] \
+--steps 50 --sr_scale 1 \
+--input [img_dir_path] \
+--color_fix_type wavelet \
+--output [output_dir_path] \
+--disable_preprocess_model \
+--device cuda
+```
+
+## <a name="train"></a>Train
+
+### Degradation Details
+
+For general image restoration, we first train both the stage1 and stage2 model under codeformer degradation to enhance the generative capacity of the stage2 model. In order to improve the ability for degradation removal, we train another stage1 model under Real-ESRGAN degradation and utilize it during inference.
+
+For face image restoration, we adopt the degradation model used in [DifFace](https://github.com/zsyOAOA/DifFace/blob/master/configs/training/swinir_ffhq512.yaml) for training and directly utilize the SwinIR model released by them as our stage1 model.
+
+### Data Preparation
+
+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_codeformer_train.yaml) and [validation set](configs/dataset/general_deg_codeformer_val.yaml) for **CodeFormer** degradation.
+    - [training set](configs/dataset/general_deg_realesrgan_train.yaml) and [validation set](configs/dataset/general_deg_realesrgan_val.yaml) for **Real-ESRGAN** degradation.
+
+    For face image restoration, fill in the face [training set](configs/dataset/face_train.yaml) and [validation set](configs/dataset/face_val.yaml) configuration files with appropriate values.
+
+### Train Stage1 Model
+
+1. Configure training-related information.
+
+    Fill in the configuration file of [training](configs/train_swinir.yaml) with appropriate values.
+
+2. Start training.
+
+    ```shell
+    python train.py --config [training_config_path]
+    ```
+
+    :bulb::Checkpoints of SwinIR will be used in training stage2 model.
+
+### Train Stage2 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/cldm.yaml \
+    --sd_weight [sd_v2.1_ckpt_path] \
+    --swinir_weight [swinir_ckpt_path] \
+    --output [init_weight_output_path]
+    ```
+
+    You will see some [outputs](assets/init_weight_outputs.txt) which show the weight initialization.
+
+3. Configure training-related information.
+
+    Fill in the configuration file of [training](configs/train_cldm.yaml) with appropriate values.
+
+4. Start training.
+
+    ```shell
+    python train.py --config [training_config_path]
+    ```
+
+## Citation
+
+Please cite us if our work is useful for your research.
+
+```
+@article{2023diffbir,
+  author    = {Xinqi Lin, Jingwen He, Ziyan Chen, Zhaoyang Lyu, Ben Fei, Bo Dai, Wanli Ouyang, Yu Qiao, Chao Dong},
+  title     = {DiffBIR: Towards Blind Image Restoration with Generative Diffusion Prior},
+  journal   = {arxiv},
+  year      = {2023},
+}
+```
+
+## License
+
+This project is released under the [Apache 2.0 license](LICENSE).
+
+## Acknowledgement
+
+This project is based on [ControlNet](https://github.com/lllyasviel/ControlNet) and [BasicSR](https://github.com/XPixelGroup/BasicSR). Thanks for their awesome work.
+
+## Contact
+
+If you have any questions, please feel free to contact with me at linxinqi@tju.edu.cn.

gitattributes.txt

@@ -0,0 +1,39 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+diffbir_face_full_v1 filter=lfs diff=lfs merge=lfs -text
+diffbir_face_swinir_v1 filter=lfs diff=lfs merge=lfs -text
+diffbir_general_full_v1 filter=lfs diff=lfs merge=lfs -text
+diffbir_general_swinir_v1 filter=lfs diff=lfs merge=lfs -text

gitignore.txt

@@ -0,0 +1,17 @@
+__pycache__
+*.ckpt
+*.pth
+/data
+/exps
+*.sh
+!install_env.sh
+/weights
+/temp
+/results
+.ipynb_checkpoints/
+/TODO.txt
+/deprecated
+/temp_scripts
+/.vscode
+/runs
+/tests

inference.py

@@ -0,0 +1,86 @@
+from argparse import ArgumentParser, Namespace
+
+import torch
+
+from accelerate.utils import set_seed
+from utils.inference import (
+    V1InferenceLoop,
+    BSRInferenceLoop, BFRInferenceLoop, BIDInferenceLoop, UnAlignedBFRInferenceLoop
+)
+
+
+def check_device(device: str) -> str:
+    if device == "cuda":
+        if not torch.cuda.is_available():
+            print("CUDA not available because the current PyTorch install was not "
+                  "built with CUDA enabled.")
+            device = "cpu"
+    else:
+        if device == "mps":
+            if not torch.backends.mps.is_available():
+                if not torch.backends.mps.is_built():
+                    print("MPS not available because the current PyTorch install was not "
+                          "built with MPS enabled.")
+                    device = "cpu"
+                else:
+                    print("MPS not available because the current MacOS version is not 12.3+ "
+                          "and/or you do not have an MPS-enabled device on this machine.")
+                    device = "cpu"
+    print(f"using device {device}")
+    return device
+
+
+def parse_args() -> Namespace:
+    parser = ArgumentParser()
+    ### model parameters
+    parser.add_argument("--task", type=str, required=True, choices=["sr", "dn", "fr", "fr_bg"])
+    parser.add_argument("--upscale", type=float, required=True)
+    parser.add_argument("--version", type=str, default="v2", choices=["v1", "v2"])
+    ### sampling parameters
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--better_start", action="store_true")
+    parser.add_argument("--tiled", action="store_true")
+    parser.add_argument("--tile_size", type=int, default=512)
+    parser.add_argument("--tile_stride", type=int, default=256)
+    parser.add_argument("--pos_prompt", type=str, default="")
+    parser.add_argument("--neg_prompt", type=str, default="low quality, blurry, low-resolution, noisy, unsharp, weird textures")
+    parser.add_argument("--cfg_scale", type=float, default=4.0)
+    ### input parameters
+    parser.add_argument("--input", type=str, required=True)
+    parser.add_argument("--n_samples", type=int, default=1)
+    ### guidance parameters
+    parser.add_argument("--guidance", action="store_true")
+    parser.add_argument("--g_loss", type=str, default="w_mse", choices=["mse", "w_mse"])
+    parser.add_argument("--g_scale", type=float, default=0.0)
+    parser.add_argument("--g_start", type=int, default=1001)
+    parser.add_argument("--g_stop", type=int, default=-1)
+    parser.add_argument("--g_space", type=str, default="latent")
+    parser.add_argument("--g_repeat", type=int, default=1)
+    ### output parameters
+    parser.add_argument("--output", type=str, required=True)
+    ### common parameters
+    parser.add_argument("--seed", type=int, default=231)
+    parser.add_argument("--device", type=str, default="cuda", choices=["cpu", "cuda", "mps"])
+    
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    args.device = check_device(args.device)
+    set_seed(args.seed)
+    if args.version == "v1":
+        V1InferenceLoop(args).run()
+    else:
+        supported_tasks = {
+            "sr": BSRInferenceLoop,
+            "dn": BIDInferenceLoop,
+            "fr": BFRInferenceLoop,
+            "fr_bg": UnAlignedBFRInferenceLoop
+        }
+        supported_tasks[args.task](args).run()
+        print("done!")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,19 @@
+--extra-index-url https://download.pytorch.org/whl/cu118
+torch==2.2.2+cu118
+torchvision==0.17.2+cu118
+torchaudio==2.2.2+cu118
+omegaconf==2.3.0
+accelerate==0.28.0
+einops==0.7.0
+opencv_python==4.9.0.80
+scipy==1.12.0
+ftfy==6.2.0
+regex==2023.12.25
+python-dateutil==2.9.0.post0
+timm==0.9.16
+pytorch-lightning==2.2.1 # only for loading pretrained sd weight
+tensorboard==2.16.2 # for tensorboard event visualization
+protobuf==4.25.3 # for tensorboard
+lpips==0.1.4
+facexlib==0.3.0
+gradio==4.25.0

train_stage1.py

@@ -0,0 +1,249 @@
+import os
+from argparse import ArgumentParser
+import warnings
+
+from omegaconf import OmegaConf
+import torch
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from torchvision.utils import make_grid
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from einops import rearrange
+from tqdm import tqdm
+import lpips
+
+from model import SwinIR
+from utils.common import instantiate_from_config
+
+
+# https://github.com/XPixelGroup/BasicSR/blob/033cd6896d898fdd3dcda32e3102a792efa1b8f4/basicsr/utils/color_util.py#L186
+def rgb2ycbcr_pt(img, y_only=False):
+    """Convert RGB images to YCbCr images (PyTorch version).
+
+    It implements the ITU-R BT.601 conversion for standard-definition television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+
+    Args:
+        img (Tensor): Images with shape (n, 3, h, w), the range [0, 1], float, RGB format.
+         y_only (bool): Whether to only return Y channel. Default: False.
+
+    Returns:
+        (Tensor): converted images with the shape (n, 3/1, h, w), the range [0, 1], float.
+    """
+    if y_only:
+        weight = torch.tensor([[65.481], [128.553], [24.966]]).to(img)
+        out_img = torch.matmul(img.permute(0, 2, 3, 1), weight).permute(0, 3, 1, 2) + 16.0
+    else:
+        weight = torch.tensor([[65.481, -37.797, 112.0], [128.553, -74.203, -93.786], [24.966, 112.0, -18.214]]).to(img)
+        bias = torch.tensor([16, 128, 128]).view(1, 3, 1, 1).to(img)
+        out_img = torch.matmul(img.permute(0, 2, 3, 1), weight).permute(0, 3, 1, 2) + bias
+
+    out_img = out_img / 255.
+    return out_img
+
+
+# https://github.com/XPixelGroup/BasicSR/blob/033cd6896d898fdd3dcda32e3102a792efa1b8f4/basicsr/metrics/psnr_ssim.py#L52
+def calculate_psnr_pt(img, img2, crop_border, test_y_channel=False):
+    """Calculate PSNR (Peak Signal-to-Noise Ratio) (PyTorch version).
+
+    Reference: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
+
+    Args:
+        img (Tensor): Images with range [0, 1], shape (n, 3/1, h, w).
+        img2 (Tensor): Images with range [0, 1], shape (n, 3/1, h, w).
+        crop_border (int): Cropped pixels in each edge of an image. These pixels are not involved in the calculation.
+        test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
+
+    Returns:
+        float: PSNR result.
+    """
+
+    assert img.shape == img2.shape, (f'Image shapes are different: {img.shape}, {img2.shape}.')
+
+    if crop_border != 0:
+        img = img[:, :, crop_border:-crop_border, crop_border:-crop_border]
+        img2 = img2[:, :, crop_border:-crop_border, crop_border:-crop_border]
+
+    if test_y_channel:
+        img = rgb2ycbcr_pt(img, y_only=True)
+        img2 = rgb2ycbcr_pt(img2, y_only=True)
+
+    img = img.to(torch.float64)
+    img2 = img2.to(torch.float64)
+
+    mse = torch.mean((img - img2)**2, dim=[1, 2, 3])
+    return 10. * torch.log10(1. / (mse + 1e-8))
+
+
+def main(args) -> None:
+    # Setup accelerator:
+    accelerator = Accelerator(split_batches=True)
+    set_seed(231)
+    device = accelerator.device
+    cfg = OmegaConf.load(args.config)
+
+    # Setup an experiment folder:
+    if accelerator.is_local_main_process:
+        exp_dir = cfg.train.exp_dir
+        os.makedirs(exp_dir, exist_ok=True)
+        ckpt_dir = os.path.join(exp_dir, "checkpoints")
+        os.makedirs(ckpt_dir, exist_ok=True)
+        print(f"Experiment directory created at {exp_dir}")
+
+    # Create model:
+    swinir: SwinIR = instantiate_from_config(cfg.model.swinir)
+    if cfg.train.resume:
+        swinir.load_state_dict(torch.load(cfg.train.resume, map_location="cpu"), strict=True)
+        if accelerator.is_local_main_process:
+            print(f"strictly load weight from checkpoint: {cfg.train.resume}")
+    else:
+        if accelerator.is_local_main_process:
+            print("initialize from scratch")
+    
+    # Setup optimizer:
+    opt = torch.optim.AdamW(
+        swinir.parameters(), lr=cfg.train.learning_rate,
+        weight_decay=0
+    )
+    
+    # Setup data:
+    dataset = instantiate_from_config(cfg.dataset.train)
+    loader = DataLoader(
+        dataset=dataset, batch_size=cfg.train.batch_size,
+        num_workers=cfg.train.num_workers,
+        shuffle=True, drop_last=True
+    )
+    val_dataset = instantiate_from_config(cfg.dataset.val)
+    val_loader = DataLoader(
+        dataset=val_dataset, batch_size=cfg.train.batch_size,
+        num_workers=cfg.train.num_workers,
+        shuffle=False, drop_last=False
+    )
+    if accelerator.is_local_main_process:
+        print(f"Dataset contains {len(dataset):,} images from {dataset.file_list}")
+
+    # Prepare models for training:
+    swinir.train().to(device)
+    swinir, opt, loader, val_loader = accelerator.prepare(swinir, opt, loader, val_loader)
+    pure_swinir = accelerator.unwrap_model(swinir)
+
+    # Variables for monitoring/logging purposes:
+    global_step = 0
+    max_steps = cfg.train.train_steps
+    step_loss = []
+    epoch = 0
+    epoch_loss = []
+    with warnings.catch_warnings():
+        # avoid warnings from lpips internal
+        warnings.simplefilter("ignore")
+        lpips_model = lpips.LPIPS(net="alex", verbose=accelerator.is_local_main_process).eval().to(device)
+    if accelerator.is_local_main_process:
+        writer = SummaryWriter(exp_dir)
+        print(f"Training for {max_steps} steps...")
+    
+    while global_step < max_steps:
+        pbar = tqdm(iterable=None, disable=not accelerator.is_local_main_process, unit="batch", total=len(loader))
+        for gt, lq, _ in loader:
+            gt = rearrange((gt + 1) / 2, "b h w c -> b c h w").contiguous().float().to(device)
+            lq = rearrange(lq, "b h w c -> b c h w").contiguous().float().to(device)
+            pred = swinir(lq)
+            loss = F.mse_loss(input=pred, target=gt, reduction="sum")
+
+            opt.zero_grad()
+            accelerator.backward(loss)
+            opt.step()
+            accelerator.wait_for_everyone()
+
+            global_step += 1
+            step_loss.append(loss.item())
+            epoch_loss.append(loss.item())
+            pbar.update(1)
+            pbar.set_description(f"Epoch: {epoch:04d}, Global Step: {global_step:07d}, Loss: {loss.item():.6f}")
+
+            # Log loss values:
+            if global_step % cfg.train.log_every == 0:
+                # Gather values from all processes
+                avg_loss = accelerator.gather(torch.tensor(step_loss, device=device).unsqueeze(0)).mean().item()
+                step_loss.clear()
+                if accelerator.is_local_main_process:
+                    writer.add_scalar("train/loss_step", avg_loss, global_step)
+
+            # Save checkpoint:
+            if global_step % cfg.train.ckpt_every == 0:
+                if accelerator.is_local_main_process:
+                    checkpoint = pure_swinir.state_dict()
+                    ckpt_path = f"{ckpt_dir}/{global_step:07d}.pt"
+                    torch.save(checkpoint, ckpt_path)
+
+            if global_step % cfg.train.image_every == 0 or global_step == 1:
+                swinir.eval()
+                N = 12
+                log_gt, log_lq = gt[:N], lq[:N]
+                with torch.no_grad():
+                    log_pred = swinir(log_lq)
+                if accelerator.is_local_main_process:
+                    for tag, image in [
+                        ("image/pred", log_pred),
+                        ("image/gt", log_gt),
+                        ("image/lq", log_lq),
+                    ]:
+                        writer.add_image(tag, make_grid(image, nrow=4), global_step)
+                swinir.train()
+
+            # Evaluate model:
+            if global_step % cfg.train.val_every == 0:
+                swinir.eval()
+                val_loss = []
+                val_lpips = []
+                val_psnr = []
+                val_pbar = tqdm(iterable=None, disable=not accelerator.is_local_main_process, unit="batch",
+                                total=len(val_loader), leave=False, desc="Validation")
+                # TODO: use accelerator.gather_for_metrics for more precise metric calculation?
+                for val_gt, val_lq, _ in val_loader:
+                    val_gt = rearrange((val_gt + 1) / 2, "b h w c -> b c h w").contiguous().float().to(device)
+                    val_lq = rearrange(val_lq, "b h w c -> b c h w").contiguous().float().to(device)
+                    with torch.no_grad():
+                        # forward
+                        val_pred = swinir(val_lq)
+                        # compute metrics (loss, lpips, psnr)
+                        val_loss.append(F.mse_loss(input=val_pred, target=val_gt, reduction="sum").item())
+                        val_lpips.append(lpips_model(val_pred, val_gt, normalize=True).mean().item())
+                        val_psnr.append(calculate_psnr_pt(val_pred, val_gt, crop_border=0).mean().item())
+                    val_pbar.update(1)
+                val_pbar.close()
+                avg_val_loss = accelerator.gather(torch.tensor(val_loss, device=device).unsqueeze(0)).mean().item()
+                avg_val_lpips = accelerator.gather(torch.tensor(val_lpips, device=device).unsqueeze(0)).mean().item()
+                avg_val_psnr = accelerator.gather(torch.tensor(val_psnr, device=device).unsqueeze(0)).mean().item()
+                if accelerator.is_local_main_process:
+                    for tag, val in [
+                        ("val/loss", avg_val_loss),
+                        ("val/lpips", avg_val_lpips),
+                        ("val/psnr", avg_val_psnr)
+                    ]:
+                        writer.add_scalar(tag, val, global_step)
+                swinir.train()
+            
+            accelerator.wait_for_everyone()
+
+            if global_step == max_steps:
+                break
+        
+        pbar.close()
+        epoch += 1
+        avg_epoch_loss = accelerator.gather(torch.tensor(epoch_loss, device=device).unsqueeze(0)).mean().item()
+        epoch_loss.clear()
+        if accelerator.is_local_main_process:
+            writer.add_scalar("train/loss_epoch", avg_epoch_loss, global_step)
+
+    if accelerator.is_local_main_process:
+        print("done!")
+        writer.close()
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--config", type=str, required=True)
+    args = parser.parse_args()
+    main(args)

train_stage2.py

@@ -0,0 +1,207 @@
+import os
+from argparse import ArgumentParser
+
+from omegaconf import OmegaConf
+import torch
+from torch.utils.data import DataLoader
+from torchvision.utils import make_grid
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+from einops import rearrange
+from tqdm import tqdm
+from torch.utils.tensorboard import SummaryWriter
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+
+from model import ControlLDM, SwinIR, Diffusion
+from utils.common import instantiate_from_config
+from utils.sampler import SpacedSampler
+
+
+def log_txt_as_img(wh, xc):
+    # wh a tuple of (width, height)
+    # xc a list of captions to plot
+    b = len(xc)
+    txts = list()
+    for bi in range(b):
+        txt = Image.new("RGB", wh, color="white")
+        draw = ImageDraw.Draw(txt)
+        # font = ImageFont.truetype('font/DejaVuSans.ttf', size=size)
+        font = ImageFont.load_default()
+        nc = int(40 * (wh[0] / 256))
+        lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc))
+
+        try:
+            draw.text((0, 0), lines, fill="black", font=font)
+        except UnicodeEncodeError:
+            print("Cant encode string for logging. Skipping.")
+
+        txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0
+        txts.append(txt)
+    txts = np.stack(txts)
+    txts = torch.tensor(txts)
+    return txts
+
+
+def main(args) -> None:
+    # Setup accelerator:
+    accelerator = Accelerator(split_batches=True)
+    set_seed(231)
+    device = accelerator.device
+    cfg = OmegaConf.load(args.config)
+
+    # Setup an experiment folder:
+    if accelerator.is_local_main_process:
+        exp_dir = cfg.train.exp_dir
+        os.makedirs(exp_dir, exist_ok=True)
+        ckpt_dir = os.path.join(exp_dir, "checkpoints")
+        os.makedirs(ckpt_dir, exist_ok=True)
+        print(f"Experiment directory created at {exp_dir}")
+
+    # Create model:
+    cldm: ControlLDM = instantiate_from_config(cfg.model.cldm)
+    sd = torch.load(cfg.train.sd_path, map_location="cpu")["state_dict"]
+    unused = cldm.load_pretrained_sd(sd)
+    if accelerator.is_local_main_process:
+        print(f"strictly load pretrained SD weight from {cfg.train.sd_path}\n"
+              f"unused weights: {unused}")
+    
+    if cfg.train.resume:
+        cldm.load_controlnet_from_ckpt(torch.load(cfg.train.resume, map_location="cpu"))
+        if accelerator.is_local_main_process:
+            print(f"strictly load controlnet weight from checkpoint: {cfg.train.resume}")
+    else:
+        init_with_new_zero, init_with_scratch = cldm.load_controlnet_from_unet()
+        if accelerator.is_local_main_process:
+            print(f"strictly load controlnet weight from pretrained SD\n"
+                  f"weights initialized with newly added zeros: {init_with_new_zero}\n"
+                  f"weights initialized from scratch: {init_with_scratch}")
+    
+    swinir: SwinIR = instantiate_from_config(cfg.model.swinir)
+    sd = {
+        (k[len("module."):] if k.startswith("module.") else k): v
+        for k, v in torch.load(cfg.train.swinir_path, map_location="cpu").items()
+    }
+    swinir.load_state_dict(sd, strict=True)
+    for p in swinir.parameters():
+        p.requires_grad = False
+    if accelerator.is_local_main_process:
+        print(f"load SwinIR from {cfg.train.swinir_path}")
+    
+    diffusion: Diffusion = instantiate_from_config(cfg.model.diffusion)
+    
+    # Setup optimizer:
+    opt = torch.optim.AdamW(cldm.controlnet.parameters(), lr=cfg.train.learning_rate)
+    
+    # Setup data:
+    dataset = instantiate_from_config(cfg.dataset.train)
+    loader = DataLoader(
+        dataset=dataset, batch_size=cfg.train.batch_size,
+        num_workers=cfg.train.num_workers,
+        shuffle=True, drop_last=True
+    )
+    if accelerator.is_local_main_process:
+        print(f"Dataset contains {len(dataset):,} images from {dataset.file_list}")
+
+    # Prepare models for training:
+    cldm.train().to(device)
+    swinir.eval().to(device)
+    diffusion.to(device)
+    cldm, opt, loader = accelerator.prepare(cldm, opt, loader)
+    pure_cldm: ControlLDM = accelerator.unwrap_model(cldm)
+    
+    # Variables for monitoring/logging purposes:
+    global_step = 0
+    max_steps = cfg.train.train_steps
+    step_loss = []
+    epoch = 0
+    epoch_loss = []
+    sampler = SpacedSampler(diffusion.betas)
+    if accelerator.is_local_main_process:
+        writer = SummaryWriter(exp_dir)
+        print(f"Training for {max_steps} steps...")
+    
+    while global_step < max_steps:
+        pbar = tqdm(iterable=None, disable=not accelerator.is_local_main_process, unit="batch", total=len(loader))
+        for gt, lq, prompt in loader:
+            gt = rearrange(gt, "b h w c -> b c h w").contiguous().float().to(device)
+            lq = rearrange(lq, "b h w c -> b c h w").contiguous().float().to(device)
+            with torch.no_grad():
+                z_0 = pure_cldm.vae_encode(gt)
+                clean = swinir(lq)
+                cond = pure_cldm.prepare_condition(clean, prompt)
+            t = torch.randint(0, diffusion.num_timesteps, (z_0.shape[0],), device=device)
+            
+            loss = diffusion.p_losses(cldm, z_0, t, cond)
+            opt.zero_grad()
+            accelerator.backward(loss)
+            opt.step()
+
+            accelerator.wait_for_everyone()
+
+            global_step += 1
+            step_loss.append(loss.item())
+            epoch_loss.append(loss.item())
+            pbar.update(1)
+            pbar.set_description(f"Epoch: {epoch:04d}, Global Step: {global_step:07d}, Loss: {loss.item():.6f}")
+
+            # Log loss values:
+            if global_step % cfg.train.log_every == 0 and global_step > 0:
+                # Gather values from all processes
+                avg_loss = accelerator.gather(torch.tensor(step_loss, device=device).unsqueeze(0)).mean().item()
+                step_loss.clear()
+                if accelerator.is_local_main_process:
+                    writer.add_scalar("loss/loss_simple_step", avg_loss, global_step)
+
+            # Save checkpoint:
+            if global_step % cfg.train.ckpt_every == 0 and global_step > 0:
+                if accelerator.is_local_main_process:
+                    checkpoint = pure_cldm.controlnet.state_dict()
+                    ckpt_path = f"{ckpt_dir}/{global_step:07d}.pt"
+                    torch.save(checkpoint, ckpt_path)
+
+            if global_step % cfg.train.image_every == 0 or global_step == 1:
+                N = 12
+                log_clean = clean[:N]
+                log_cond = {k:v[:N] for k, v in cond.items()}
+                log_gt, log_lq = gt[:N], lq[:N]
+                log_prompt = prompt[:N]
+                cldm.eval()
+                with torch.no_grad():
+                    z = sampler.sample(
+                        model=cldm, device=device, steps=50, batch_size=len(log_gt), x_size=z_0.shape[1:],
+                        cond=log_cond, uncond=None, cfg_scale=1.0, x_T=None,
+                        progress=accelerator.is_local_main_process, progress_leave=False
+                    )
+                    if accelerator.is_local_main_process:
+                        for tag, image in [
+                            ("image/samples", (pure_cldm.vae_decode(z) + 1) / 2),
+                            ("image/gt", (log_gt + 1) / 2),
+                            ("image/lq", log_lq),
+                            ("image/condition", log_clean),
+                            ("image/condition_decoded", (pure_cldm.vae_decode(log_cond["c_img"]) + 1) / 2),
+                            ("image/prompt", (log_txt_as_img((512, 512), log_prompt) + 1) / 2)
+                        ]:
+                            writer.add_image(tag, make_grid(image, nrow=4), global_step)
+                cldm.train()
+            accelerator.wait_for_everyone()
+            if global_step == max_steps:
+                break
+        
+        pbar.close()
+        epoch += 1
+        avg_epoch_loss = accelerator.gather(torch.tensor(epoch_loss, device=device).unsqueeze(0)).mean().item()
+        epoch_loss.clear()
+        if accelerator.is_local_main_process:
+            writer.add_scalar("loss/loss_simple_epoch", avg_epoch_loss, global_step)
+
+    if accelerator.is_local_main_process:
+        print("done!")
+        writer.close()
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--config", type=str, required=True)
+    args = parser.parse_args()
+    main(args)