Skip to content

GitLab

"references/git@developer.sourcefind.cn:OpenDAS/vision.git" did not exist on "b37c8a3ca4c9c626cdac763c6be697231665b0f8"
Commit 48ef4486 authored by zycXD's avatar zycXD
Browse files

fix bugs

parents 705e02f6 30355a12
Hide whitespace changes
Inline Side-by-side
Showing with 434 additions and 86 deletions
.gitignore
View file @ 48ef4486
...@@ -7,4 +7,4 @@ __pycache__ ...@@ -7,4 +7,4 @@ __pycache__
!install_env.sh !install_env.sh
/weights /weights
/temp /temp
results/ /results
README.md 0 → 100644
View file @ 48ef4486
<p align="center">
<img src="assets/logo.png" width="400">
</p>
## 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 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
<p align="center">
<img src="assets/architecture.png" style="border-radius: 15px">
</p>
:star:If DiffBIR is helpful for you, please help star this repo. Thanks!:hugs:
## :book:Table Of Contents
- [Visual Results On Real-world Images](#visual_results)
- [Installation](#installation)
- [Pretrained Models](#pretrained_models)
- [Quick Start (gradio demo)](#quick_start)
- [Inference](#inference)
- [Train](#train)
- [Update](#update)
- [TODO](#todo)
## <a name="visual_results"></a>:eyes:Visual Results On Real-world Images
<!-- <details close>
<summary>General Image Restoration</summary> -->
### General Image Restoration
[<img src="assets/visual_results/general6.png" height="223px"/>](https://imgsli.com/MTk5ODI3) [<img src="assets/visual_results/general7.png" height="223px"/>](https://imgsli.com/MTk5ODI4) [<img src="assets/visual_results/general4.png" height="223px"/>](https://imgsli.com/MTk5ODI1)
[<img src="assets/visual_results/general1.png" height="223px"/>](https://imgsli.com/MTk5ODIy) [<img src="assets/visual_results/general2.png" height="223px"/>](https://imgsli.com/MTk5ODIz)
[<img src="assets/visual_results/general3.png" height="223px"/>](https://imgsli.com/MTk5ODI0) [<img src="assets/visual_results/general5.png" height="223px"/>](https://imgsli.com/MjAxMjM0)
<!-- </details> -->
<!-- <details close> -->
<!-- <summary>Face Image Restoration</summary> -->
### Face Image Restoration
[<img src="assets/visual_results/face1.png" height="223px"/>](https://imgsli.com/MTk5ODI5) [<img src="assets/visual_results/face2.png" height="223px"/>](https://imgsli.com/MTk5ODMw) [<img src="assets/visual_results/face3.png" height="223px"/>](https://imgsli.com/MTk5ODMy)
[<img src="assets/visual_results/face4.png" height="223px"/>](https://imgsli.com/MTk5ODM0) [<img src="assets/visual_results/face5.png" height="223px"/>](https://imgsli.com/MTk5ODM1) [<img src="assets/visual_results/face6.png" height="223px"/>](https://imgsli.com/MTk5ODM2)
[<img src="assets/visual_results/whole_image1.png" height="410px"/>](https://imgsli.com/MjA0MzQw)
<!-- </details> -->
## <a name="installation"></a>:gear: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 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>:dna:Pretrained Models
| Model Name | Description | HuggingFace | BaiduNetdisk |
| :--------- | :---------- | :---------- | :---------- |
| 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) |
| 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) |
| 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) |
| 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) |
## <a name="quick_start"></a>:flight_departure: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.
```
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
```
<div align="center">
<kbd><img src="assets/gradio.png"></img></kbd>
</div>
## <a name="inference"></a>:crossed_swords:Inference
### Full Pipeline (Remove Degradations & Refine Details)
#### 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/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 \
--image_size 512 \
--color_fix_type wavelet --resize_back \
--output results/general \
--device cuda
```
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.
```shell
# for aligned face inputs
python inference_face.py \
--config configs/model/cldm.yaml \
--ckpt weights/face_full_v1.ckpt \
--input inputs/face/aligned \
--steps 50 \
--sr_scale 1 \
--image_size 512 \
--color_fix_type wavelet \
--output results/face/aligned --resize_back \
--has_aligned \
--device cuda
# for unaligned face inputs
python inference_face.py \
--config configs/model/cldm.yaml \
--ckpt weights/face_full_v1.ckpt \
--input inputs/face/whole_img \
--steps 50 \
--sr_scale 1 \
--image_size 512 \
--color_fix_type wavelet \
--output results/face/whole_img --resize_back \
--device cuda
```
### 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 --image_size 512 \
--input [img_dir_path] \
--color_fix_type wavelet --resize_back \
--output [output_dir_path] \
--disable_preprocess_model \
--device cuda
```
## <a name="train"></a>:stars: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]
```
## <a name="update"></a>:new:Update
- **2023.08.30**: Repo is released.
- **2023.09.06**: 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.09.08**: Add support for restoring unaligned faces.
## <a name="todo"></a>:climbing:TODO
- [x] Release code and pretrained models:computer:.
- [x] Update links to paper and project page:link:.
- [ ] Release real47 testset:minidisc:.
- [ ] Reduce the memory usage of DiffBIR:smiley_cat:.
- [ ] Provide HuggingFace demo:notebook:.
- [ ] Upload inference code of latent image guidance:page_facing_up:.
- [ ] Improve the performance:superhero:.
- [ ] Add a patch-based sampling schedule:mag:.
## 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.
gradio_diffbir.py
View file @ 48ef4486
...@@ -10,6 +10,7 @@ import gradio as gr ...@@ -10,6 +10,7 @@ import gradio as gr
from PIL import Image from PIL import Image
from omegaconf import OmegaConf from omegaconf import OmegaConf
from ldm.xformers_state import disable_xformers
from model.spaced_sampler import SpacedSampler from model.spaced_sampler import SpacedSampler
from model.cldm import ControlLDM from model.cldm import ControlLDM
from utils.image import ( from utils.image import (
...@@ -23,10 +24,12 @@ parser.add_argument("--config", required=True, type=str) ...@@ -23,10 +24,12 @@ parser.add_argument("--config", required=True, type=str)
parser.add_argument("--ckpt", type=str, required=True) parser.add_argument("--ckpt", type=str, required=True)
parser.add_argument("--reload_swinir", action="store_true") parser.add_argument("--reload_swinir", action="store_true")
parser.add_argument("--swinir_ckpt", type=str, default="") parser.add_argument("--swinir_ckpt", type=str, default="")
parser.add_argument("--device", type=str, default="cuda", choices=["cpu", "cuda"])
args = parser.parse_args() args = parser.parse_args()
# load model # load model
device = "cuda" if torch.cuda.is_available() else "cpu" if args.device == "cpu":
disable_xformers()
model: ControlLDM = instantiate_from_config(OmegaConf.load(args.config)) model: ControlLDM = instantiate_from_config(OmegaConf.load(args.config))
load_state_dict(model, torch.load(args.ckpt, map_location="cpu"), strict=True) load_state_dict(model, torch.load(args.ckpt, map_location="cpu"), strict=True)
# reload preprocess model if specified # reload preprocess model if specified
...@@ -34,7 +37,7 @@ if args.reload_swinir: ...@@ -34,7 +37,7 @@ if args.reload_swinir:
print(f"reload swinir model from {args.swinir_ckpt}") print(f"reload swinir model from {args.swinir_ckpt}")
load_state_dict(model.preprocess_model, torch.load(args.swinir_ckpt, map_location="cpu"), strict=True) load_state_dict(model.preprocess_model, torch.load(args.swinir_ckpt, map_location="cpu"), strict=True)
model.freeze() model.freeze()
model.to(device) model.to(args.device)
# load sampler # load sampler
sampler = SpacedSampler(model, var_type="fixed_small") sampler = SpacedSampler(model, var_type="fixed_small")
......
inference.py
View file @ 48ef4486
...@@ -10,6 +10,7 @@ import pytorch_lightning as pl ...@@ -10,6 +10,7 @@ import pytorch_lightning as pl
from PIL import Image from PIL import Image
from omegaconf import OmegaConf from omegaconf import OmegaConf
from ldm.xformers_state import disable_xformers
from model.spaced_sampler import SpacedSampler from model.spaced_sampler import SpacedSampler
from model.ddim_sampler import DDIMSampler from model.ddim_sampler import DDIMSampler
from model.cldm import ControlLDM from model.cldm import ControlLDM
...@@ -127,6 +128,7 @@ def parse_args() -> Namespace: ...@@ -127,6 +128,7 @@ def parse_args() -> Namespace:
parser.add_argument("--skip_if_exist", action="store_true") parser.add_argument("--skip_if_exist", action="store_true")
parser.add_argument("--seed", type=int, default=231) parser.add_argument("--seed", type=int, default=231)
parser.add_argument("--device", type=str, default="cuda", choices=["cpu", "cuda"])
return parser.parse_args() return parser.parse_args()
...@@ -134,7 +136,9 @@ def parse_args() -> Namespace: ...@@ -134,7 +136,9 @@ def parse_args() -> Namespace:
def main() -> None: def main() -> None:
args = parse_args() args = parse_args()
pl.seed_everything(args.seed) pl.seed_everything(args.seed)
device = "cuda" if torch.cuda.is_available() else "cpu"
if args.device == "cpu":
disable_xformers()
model: ControlLDM = instantiate_from_config(OmegaConf.load(args.config)) model: ControlLDM = instantiate_from_config(OmegaConf.load(args.config))
load_state_dict(model, torch.load(args.ckpt, map_location="cpu"), strict=True) load_state_dict(model, torch.load(args.ckpt, map_location="cpu"), strict=True)
...@@ -145,68 +149,68 @@ def main() -> None: ...@@ -145,68 +149,68 @@ def main() -> None:
print(f"reload swinir model from {args.swinir_ckpt}") print(f"reload swinir model from {args.swinir_ckpt}")
load_state_dict(model.preprocess_model, torch.load(args.swinir_ckpt, map_location="cpu"), strict=True) load_state_dict(model.preprocess_model, torch.load(args.swinir_ckpt, map_location="cpu"), strict=True)
model.freeze() model.freeze()
model.to(device) model.to(args.device)
assert os.path.isdir(args.input) assert os.path.isdir(args.input)
print(f"sampling {args.steps} steps using {args.sampler} sampler") print(f"sampling {args.steps} steps using {args.sampler} sampler")
with torch.autocast(device): # with torch.autocast(device, dtype=torch.bfloat16):
for file_path in list_image_files(args.input, follow_links=True): for file_path in list_image_files(args.input, follow_links=True):
lq = Image.open(file_path).convert("RGB") lq = Image.open(file_path).convert("RGB")
if args.sr_scale != 1: if args.sr_scale != 1:
lq = lq.resize( lq = lq.resize(
tuple(math.ceil(x * args.sr_scale) for x in lq.size), tuple(math.ceil(x * args.sr_scale) for x in lq.size),
Image.BICUBIC Image.BICUBIC
) )
lq_resized = auto_resize(lq, args.image_size) lq_resized = auto_resize(lq, args.image_size)
x = pad(np.array(lq_resized), scale=64) x = pad(np.array(lq_resized), scale=64)
for i in range(args.repeat_times): for i in range(args.repeat_times):
save_path = os.path.join(args.output, os.path.relpath(file_path, args.input)) save_path = os.path.join(args.output, os.path.relpath(file_path, args.input))
parent_path, stem, _ = get_file_name_parts(save_path) parent_path, stem, _ = get_file_name_parts(save_path)
save_path = os.path.join(parent_path, f"{stem}_{i}.png") save_path = os.path.join(parent_path, f"{stem}_{i}.png")
if os.path.exists(save_path): if os.path.exists(save_path):
if args.skip_if_exist: if args.skip_if_exist:
print(f"skip {save_path}") print(f"skip {save_path}")
continue
else:
raise RuntimeError(f"{save_path} already exist")
os.makedirs(parent_path, exist_ok=True)
try:
preds, stage1_preds = process(
model, [x], steps=args.steps, sampler=args.sampler,
strength=1,
color_fix_type=args.color_fix_type,
disable_preprocess_model=args.disable_preprocess_model
)
except RuntimeError as e:
# Avoid cuda_out_of_memory error.
print(f"{file_path}, error: {e}")
continue continue
pred, stage1_pred = preds[0], stage1_preds[0]
# remove padding
pred = pred[:lq_resized.height, :lq_resized.width, :]
stage1_pred = stage1_pred[:lq_resized.height, :lq_resized.width, :]
if args.show_lq:
if args.resize_back:
if lq_resized.size != lq.size:
pred = np.array(Image.fromarray(pred).resize(lq.size, Image.LANCZOS))
stage1_pred = np.array(Image.fromarray(stage1_pred).resize(lq.size, Image.LANCZOS))
lq = np.array(lq)
else:
lq = np.array(lq_resized)
images = [lq, pred] if args.disable_preprocess_model else [lq, stage1_pred, pred]
Image.fromarray(np.concatenate(images, axis=1)).save(save_path)
else: else:
if args.resize_back and lq_resized.size != lq.size: raise RuntimeError(f"{save_path} already exist")
Image.fromarray(pred).resize(lq.size, Image.LANCZOS).save(save_path) os.makedirs(parent_path, exist_ok=True)
else:
Image.fromarray(pred).save(save_path) # try:
print(f"save to {save_path}") preds, stage1_preds = process(
model, [x], steps=args.steps, sampler=args.sampler,
strength=1,
color_fix_type=args.color_fix_type,
disable_preprocess_model=args.disable_preprocess_model
)
# except RuntimeError as e:
# # Avoid cuda_out_of_memory error.
# print(f"{file_path}, error: {e}")
# continue
pred, stage1_pred = preds[0], stage1_preds[0]
# remove padding
pred = pred[:lq_resized.height, :lq_resized.width, :]
stage1_pred = stage1_pred[:lq_resized.height, :lq_resized.width, :]
if args.show_lq:
if args.resize_back:
if lq_resized.size != lq.size:
pred = np.array(Image.fromarray(pred).resize(lq.size, Image.LANCZOS))
stage1_pred = np.array(Image.fromarray(stage1_pred).resize(lq.size, Image.LANCZOS))
lq = np.array(lq)
else:
lq = np.array(lq_resized)
images = [lq, pred] if args.disable_preprocess_model else [lq, stage1_pred, pred]
Image.fromarray(np.concatenate(images, axis=1)).save(save_path)
else:
if args.resize_back and lq_resized.size != lq.size:
Image.fromarray(pred).resize(lq.size, Image.LANCZOS).save(save_path)
else:
Image.fromarray(pred).save(save_path)
print(f"save to {save_path}")
if __name__ == "__main__": if __name__ == "__main__":
main() main()
inference_face.py
View file @ 48ef4486
...@@ -10,6 +10,7 @@ from argparse import ArgumentParser, Namespace ...@@ -10,6 +10,7 @@ from argparse import ArgumentParser, Namespace
from facexlib.utils.face_restoration_helper import FaceRestoreHelper from facexlib.utils.face_restoration_helper import FaceRestoreHelper
from ldm.xformers_state import disable_xformers
from model.cldm import ControlLDM from model.cldm import ControlLDM
from model.ddim_sampler import DDIMSampler from model.ddim_sampler import DDIMSampler
from model.spaced_sampler import SpacedSampler from model.spaced_sampler import SpacedSampler
...@@ -56,6 +57,7 @@ def parse_args() -> Namespace: ...@@ -56,6 +57,7 @@ def parse_args() -> Namespace:
parser.add_argument("--skip_if_exist", action="store_true") parser.add_argument("--skip_if_exist", action="store_true")
parser.add_argument("--seed", type=int, default=231) parser.add_argument("--seed", type=int, default=231)
parser.add_argument("--device", type=str, default="cuda", choices=["cpu", "cuda"])
return parser.parse_args() return parser.parse_args()
...@@ -64,7 +66,9 @@ def main() -> None: ...@@ -64,7 +66,9 @@ def main() -> None:
args = parse_args() args = parse_args()
img_save_ext = 'png' img_save_ext = 'png'
pl.seed_everything(args.seed) pl.seed_everything(args.seed)
device = "cuda" if torch.cuda.is_available() else "cpu"
if args.device == "cpu":
disable_xformers()
model: ControlLDM = instantiate_from_config(OmegaConf.load(args.config)) model: ControlLDM = instantiate_from_config(OmegaConf.load(args.config))
load_state_dict(model, torch.load(args.ckpt, map_location="cpu"), strict=True) load_state_dict(model, torch.load(args.ckpt, map_location="cpu"), strict=True)
...@@ -75,13 +79,13 @@ def main() -> None: ...@@ -75,13 +79,13 @@ def main() -> None:
print(f"reload swinir model from {args.swinir_ckpt}") print(f"reload swinir model from {args.swinir_ckpt}")
load_state_dict(model.preprocess_model, torch.load(args.swinir_ckpt, map_location="cpu"), strict=True) load_state_dict(model.preprocess_model, torch.load(args.swinir_ckpt, map_location="cpu"), strict=True)
model.freeze() model.freeze()
model.to(device) model.to(args.device)
assert os.path.isdir(args.input) assert os.path.isdir(args.input)
# ------------------ set up FaceRestoreHelper ------------------- # ------------------ set up FaceRestoreHelper -------------------
face_helper = FaceRestoreHelper( face_helper = FaceRestoreHelper(
device=device, device=args.device,
upscale_factor=1, upscale_factor=1,
face_size=args.image_size, face_size=args.image_size,
use_parse=True, use_parse=True,
...@@ -186,4 +190,4 @@ def main() -> None: ...@@ -186,4 +190,4 @@ def main() -> None:
if __name__ == "__main__": if __name__ == "__main__":
main() main()
\ No newline at end of file
ldm/modules/attention.py
View file @ 48ef4486
...@@ -7,14 +7,14 @@ from einops import rearrange, repeat ...@@ -7,14 +7,14 @@ from einops import rearrange, repeat
from typing import Optional, Any from typing import Optional, Any
from ldm.modules.diffusionmodules.util import checkpoint from ldm.modules.diffusionmodules.util import checkpoint
from ldm import xformers_state
# try:
try: # import xformers
import xformers # import xformers.ops
import xformers.ops # XFORMERS_IS_AVAILBLE = True
XFORMERS_IS_AVAILBLE = True # except:
except: # XFORMERS_IS_AVAILBLE = False
XFORMERS_IS_AVAILBLE = False
# CrossAttn precision handling # CrossAttn precision handling
import os import os
...@@ -172,7 +172,8 @@ class CrossAttention(nn.Module): ...@@ -172,7 +172,8 @@ class CrossAttention(nn.Module):
# force cast to fp32 to avoid overflowing # force cast to fp32 to avoid overflowing
if _ATTN_PRECISION =="fp32": if _ATTN_PRECISION =="fp32":
with torch.autocast(enabled=False, device_type = 'cuda'): # with torch.autocast(enabled=False, device_type = 'cuda'):
with torch.autocast(enabled=False, device_type=str(x.device)):
q, k = q.float(), k.float() q, k = q.float(), k.float()
sim = einsum('b i d, b j d -> b i j', q, k) * self.scale sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
else: else:
...@@ -230,7 +231,7 @@ class MemoryEfficientCrossAttention(nn.Module): ...@@ -230,7 +231,7 @@ class MemoryEfficientCrossAttention(nn.Module):
) )
# actually compute the attention, what we cannot get enough of # actually compute the attention, what we cannot get enough of
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op) out = xformers_state.xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
if exists(mask): if exists(mask):
raise NotImplementedError raise NotImplementedError
...@@ -251,7 +252,8 @@ class BasicTransformerBlock(nn.Module): ...@@ -251,7 +252,8 @@ class BasicTransformerBlock(nn.Module):
def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True, def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True,
disable_self_attn=False): disable_self_attn=False):
super().__init__() super().__init__()
attn_mode = "softmax-xformers" if XFORMERS_IS_AVAILBLE else "softmax" # attn_mode = "softmax-xformers" if XFORMERS_IS_AVAILBLE else "softmax"
attn_mode = "softmax-xformers" if xformers_state.is_xformers_available() else "softmax"
assert attn_mode in self.ATTENTION_MODES assert attn_mode in self.ATTENTION_MODES
attn_cls = self.ATTENTION_MODES[attn_mode] attn_cls = self.ATTENTION_MODES[attn_mode]
self.disable_self_attn = disable_self_attn self.disable_self_attn = disable_self_attn
......
ldm/modules/diffusionmodules/model.py
View file @ 48ef4486
...@@ -7,14 +7,16 @@ from einops import rearrange ...@@ -7,14 +7,16 @@ from einops import rearrange
from typing import Optional, Any from typing import Optional, Any
from ldm.modules.attention import MemoryEfficientCrossAttention from ldm.modules.attention import MemoryEfficientCrossAttention
from ldm import xformers_state
try:
import xformers # try:
import xformers.ops # import xformers
XFORMERS_IS_AVAILBLE = True # import xformers.ops
except: # XFORMERS_IS_AVAILBLE = True
XFORMERS_IS_AVAILBLE = False # except:
print("No module 'xformers'. Proceeding without it.") # XFORMERS_IS_AVAILBLE = False
# print("No module 'xformers'. Proceeding without it.")
def get_timestep_embedding(timesteps, embedding_dim): def get_timestep_embedding(timesteps, embedding_dim):
...@@ -255,7 +257,7 @@ class MemoryEfficientAttnBlock(nn.Module): ...@@ -255,7 +257,7 @@ class MemoryEfficientAttnBlock(nn.Module):
.contiguous(), .contiguous(),
(q, k, v), (q, k, v),
) )
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op) out = xformers_state.xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
out = ( out = (
out.unsqueeze(0) out.unsqueeze(0)
...@@ -279,7 +281,8 @@ class MemoryEfficientCrossAttentionWrapper(MemoryEfficientCrossAttention): ...@@ -279,7 +281,8 @@ class MemoryEfficientCrossAttentionWrapper(MemoryEfficientCrossAttention):
def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None): def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None):
assert attn_type in ["vanilla", "vanilla-xformers", "memory-efficient-cross-attn", "linear", "none"], f'attn_type {attn_type} unknown' assert attn_type in ["vanilla", "vanilla-xformers", "memory-efficient-cross-attn", "linear", "none"], f'attn_type {attn_type} unknown'
if XFORMERS_IS_AVAILBLE and attn_type == "vanilla": # if XFORMERS_IS_AVAILBLE and attn_type == "vanilla":
if xformers_state.is_xformers_available() and attn_type == "vanilla":
attn_type = "vanilla-xformers" attn_type = "vanilla-xformers"
print(f"making attention of type '{attn_type}' with {in_channels} in_channels") print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
if attn_type == "vanilla": if attn_type == "vanilla":
......
ldm/modules/encoders/modules.py
View file @ 48ef4486
...@@ -140,7 +140,8 @@ class FrozenOpenCLIPEmbedder(AbstractEncoder): ...@@ -140,7 +140,8 @@ class FrozenOpenCLIPEmbedder(AbstractEncoder):
"last", "last",
"penultimate" "penultimate"
] ]
def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77, # def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77,
def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", max_length=77,
freeze=True, layer="last"): freeze=True, layer="last"):
super().__init__() super().__init__()
assert layer in self.LAYERS assert layer in self.LAYERS
...@@ -148,7 +149,7 @@ class FrozenOpenCLIPEmbedder(AbstractEncoder): ...@@ -148,7 +149,7 @@ class FrozenOpenCLIPEmbedder(AbstractEncoder):
del model.visual del model.visual
self.model = model self.model = model
self.device = device # self.device = device
self.max_length = max_length self.max_length = max_length
if freeze: if freeze:
self.freeze() self.freeze()
...@@ -167,7 +168,8 @@ class FrozenOpenCLIPEmbedder(AbstractEncoder): ...@@ -167,7 +168,8 @@ class FrozenOpenCLIPEmbedder(AbstractEncoder):
def forward(self, text): def forward(self, text):
tokens = open_clip.tokenize(text) tokens = open_clip.tokenize(text)
z = self.encode_with_transformer(tokens.to(self.device)) # z = self.encode_with_transformer(tokens.to(self.device))
z = self.encode_with_transformer(tokens.to(next(self.model.parameters()).device))
return z return z
def encode_with_transformer(self, text): def encode_with_transformer(self, text):
......
ldm/xformers_state.py 0 → 100644
View file @ 48ef4486
try:
import xformers
import xformers.ops
XFORMERS_IS_AVAILBLE = True
except:
XFORMERS_IS_AVAILBLE = False
print("No module 'xformers'. Proceeding without it.")
def is_xformers_available() -> bool:
global XFORMERS_IS_AVAILBLE
return XFORMERS_IS_AVAILBLE
def disable_xformers() -> None:
print("DISABLE XFORMERS!")
global XFORMERS_IS_AVAILBLE
XFORMERS_IS_AVAILBLE = False
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment