Fabrice-TIERCELIN commited on
Commit
f12595e
1 Parent(s): 5a8a838

Upload 8 files

Browse files
Files changed (8) hide show
  1. CKPT_PTH.py +4 -0
  2. LICENSE +18 -0
  3. README.md +157 -0
  4. gradio_demo.py +314 -0
  5. gradio_demo_face.py +411 -0
  6. gradio_demo_tiled.py +339 -0
  7. requirements.txt +42 -0
  8. test.py +106 -0
CKPT_PTH.py ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ LLAVA_CLIP_PATH = '/opt/data/private/AIGC_pretrain/LLaVA1.5/clip-vit-large-patch14-336'
2
+ LLAVA_MODEL_PATH = '/opt/data/private/AIGC_pretrain/LLaVA1.5/llava-v1.5-13b'
3
+ SDXL_CLIP1_PATH = '/opt/data/private/AIGC_pretrain/clip-vit-large-patch14'
4
+ SDXL_CLIP2_CKPT_PTH = '/opt/data/private/AIGC_pretrain/CLIP-ViT-bigG-14-laion2B-39B-b160k/open_clip_pytorch_model.bin'
LICENSE ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ License
2
+
3
+ Copyright (c) 2024 XPixel Group, Especially the author team of SUPIR.
4
+
5
+ The SUPIR ("Software") is made available for use, reproduction, and distribution strictly for non-commercial purposes. For the purposes of this declaration, "non-commercial" is defined as not primarily intended for or directed towards commercial advantage or monetary compensation.
6
+
7
+ By using, reproducing, or distributing the Software, you agree to abide by this restriction and not to use the Software for any commercial purposes without obtaining prior written permission from Dr. Jinjin Gu.
8
+
9
+ This declaration does not in any way limit the rights under any open source license that may apply to the Software; it solely adds a condition that the Software shall not be used for commercial purposes.
10
+
11
+ IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
12
+
13
+ For inquiries or to obtain permission for commercial use, please contact Dr. Jinjin Gu (jinjin.gu@suppixel.ai).
14
+
15
+
16
+
17
+
18
+
README.md ADDED
@@ -0,0 +1,157 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ## (CVPR2024) Scaling Up to Excellence: Practicing Model Scaling for Photo-Realistic Image Restoration In the Wild
2
+
3
+ > [[Paper](https://arxiv.org/abs/2401.13627)] &emsp; [[Project Page](http://supir.xpixel.group/)] &emsp; [Online Demo (Coming soon)] <br>
4
+ > Fanghua, Yu, [Jinjin Gu](https://www.jasongt.com/), Zheyuan Li, Jinfan Hu, Xiangtao Kong, [Xintao Wang](https://xinntao.github.io/), [Jingwen He](https://scholar.google.com.hk/citations?user=GUxrycUAAAAJ), [Yu Qiao](https://scholar.google.com.hk/citations?user=gFtI-8QAAAAJ), [Chao Dong](https://scholar.google.com.hk/citations?user=OSDCB0UAAAAJ) <br>
5
+ > Shenzhen Institute of Advanced Technology; Shanghai AI Laboratory; University of Sydney; The Hong Kong Polytechnic University; ARC Lab, Tencent PCG; The Chinese University of Hong Kong <br>
6
+
7
+
8
+ <p align="center">
9
+ <img src="assets/teaser.png">
10
+ </p>
11
+
12
+ ---
13
+ #### ⚠ Due to the large RAM (60G) and VRAM (30G x2) costs of SUPIR, we are working on the online demo releasing.
14
+
15
+ ---
16
+ ## 🔧 Dependencies and Installation
17
+
18
+ 1. Clone repo
19
+ ```bash
20
+ git clone https://github.com/Fanghua-Yu/SUPIR.git
21
+ cd SUPIR
22
+ ```
23
+
24
+ 2. Install dependent packages
25
+ ```bash
26
+ conda create -n SUPIR python=3.8 -y
27
+ conda activate SUPIR
28
+ pip install --upgrade pip
29
+ pip install -r requirements.txt
30
+ ```
31
+
32
+ 3. Download Checkpoints
33
+
34
+ For users who can connect to huggingface, please setting `LLAVA_CLIP_PATH, SDXL_CLIP1_PATH, SDXL_CLIP2_CKPT_PTH` in `CKPT_PTH.py` as `None`. These CLIPs will be downloaded automatically.
35
+
36
+ #### Dependent Models
37
+ * [SDXL CLIP Encoder-1](https://huggingface.co/openai/clip-vit-large-patch14)
38
+ * [SDXL CLIP Encoder-2](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)
39
+ * [SDXL base 1.0_0.9vae](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors)
40
+ * [LLaVA CLIP](https://huggingface.co/openai/clip-vit-large-patch14-336)
41
+ * [LLaVA v1.5 13B](https://huggingface.co/liuhaotian/llava-v1.5-13b)
42
+ * (optional) [Juggernaut-XL_v9_RunDiffusionPhoto_v2](https://huggingface.co/RunDiffusion/Juggernaut-XL-v9/blob/main/Juggernaut-XL_v9_RunDiffusionPhoto_v2.safetensors)
43
+ * Replacement of `SDXL base 1.0_0.9vae` for Photo Realistic
44
+ * (optional) [Juggernaut_RunDiffusionPhoto2_Lightning_4Steps](https://huggingface.co/RunDiffusion/Juggernaut-XL-Lightning/blob/main/Juggernaut_RunDiffusionPhoto2_Lightning_4Steps.safetensors)
45
+ * Distilling model used in `SUPIR_v0_Juggernautv9_lightning.yaml`
46
+
47
+
48
+ #### Models we provided:
49
+ * `SUPIR-v0Q`: [Baidu Netdisk](https://pan.baidu.com/s/1lnefCZhBTeDWijqbj1jIyw?pwd=pjq6), [Google Drive](https://drive.google.com/drive/folders/1yELzm5SvAi9e7kPcO_jPp2XkTs4vK6aR?usp=sharing)
50
+
51
+ Default training settings with paper. High generalization and high image quality in most cases.
52
+
53
+ * `SUPIR-v0F`: [Baidu Netdisk](https://pan.baidu.com/s/1AECN8NjiVuE3hvO8o-Ua6A?pwd=k2uz), [Google Drive](https://drive.google.com/drive/folders/1yELzm5SvAi9e7kPcO_jPp2XkTs4vK6aR?usp=sharing)
54
+
55
+ Training with light degradation settings. Stage1 encoder of `SUPIR-v0F` remains more details when facing light degradations.
56
+
57
+ 4. Edit Custom Path for Checkpoints
58
+ ```
59
+ * [CKPT_PTH.py] --> LLAVA_CLIP_PATH, LLAVA_MODEL_PATH, SDXL_CLIP1_PATH, SDXL_CLIP2_CACHE_DIR
60
+ * [options/SUPIR_v0.yaml] --> SDXL_CKPT, SUPIR_CKPT_Q, SUPIR_CKPT_F
61
+ ```
62
+ ---
63
+
64
+ ## ⚡ Quick Inference
65
+ ### Val Dataset
66
+ RealPhoto60: [Baidu Netdisk](https://pan.baidu.com/s/1CJKsPGtyfs8QEVCQ97voBA?pwd=aocg), [Google Drive](https://drive.google.com/drive/folders/1yELzm5SvAi9e7kPcO_jPp2XkTs4vK6aR?usp=sharing)
67
+
68
+ ### Usage of SUPIR
69
+ ```Shell
70
+ Usage:
71
+ -- python test.py [options]
72
+ -- python gradio_demo.py [interactive options]
73
+
74
+ --img_dir Input folder.
75
+ --save_dir Output folder.
76
+ --upscale Upsampling ratio of given inputs. Default: 1
77
+ --SUPIR_sign Model selection. Default: 'Q'; Options: ['F', 'Q']
78
+ --seed Random seed. Default: 1234
79
+ --min_size Minimum resolution of output images. Default: 1024
80
+ --edm_steps Numb of steps for EDM Sampling Scheduler. Default: 50
81
+ --s_stage1 Control Strength of Stage1. Default: -1 (negative means invalid)
82
+ --s_churn Original hy-param of EDM. Default: 5
83
+ --s_noise Original hy-param of EDM. Default: 1.003
84
+ --s_cfg Classifier-free guidance scale for prompts. Default: 7.5
85
+ --s_stage2 Control Strength of Stage2. Default: 1.0
86
+ --num_samples Number of samples for each input. Default: 1
87
+ --a_prompt Additive positive prompt for all inputs.
88
+ Default: 'Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera,
89
+ hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme
90
+ meticulous detailing, skin pore detailing, hyper sharpness, perfect without deformations.'
91
+ --n_prompt Fixed negative prompt for all inputs.
92
+ Default: 'painting, oil painting, illustration, drawing, art, sketch, oil painting,
93
+ cartoon, CG Style, 3D render, unreal engine, blurring, dirty, messy, worst quality,
94
+ low quality, frames, watermark, signature, jpeg artifacts, deformed, lowres, over-smooth'
95
+ --color_fix_type Color Fixing Type. Default: 'Wavelet'; Options: ['None', 'AdaIn', 'Wavelet']
96
+ --linear_CFG Linearly (with sigma) increase CFG from 'spt_linear_CFG' to s_cfg. Default: False
97
+ --linear_s_stage2 Linearly (with sigma) increase s_stage2 from 'spt_linear_s_stage2' to s_stage2. Default: False
98
+ --spt_linear_CFG Start point of linearly increasing CFG. Default: 1.0
99
+ --spt_linear_s_stage2 Start point of linearly increasing s_stage2. Default: 0.0
100
+ --ae_dtype Inference data type of AutoEncoder. Default: 'bf16'; Options: ['fp32', 'bf16']
101
+ --diff_dtype Inference data type of Diffusion. Default: 'fp16'; Options: ['fp32', 'fp16', 'bf16']
102
+ ```
103
+
104
+ ### Python Script
105
+ ```Shell
106
+ # Seek for best quality for most cases
107
+ CUDA_VISIBLE_DEVICES=0,1 python test.py --img_dir '/opt/data/private/LV_Dataset/DiffGLV-Test-All/RealPhoto60/LQ' --save_dir ./results-Q --SUPIR_sign Q --upscale 2
108
+ # for light degradation and high fidelity
109
+ CUDA_VISIBLE_DEVICES=0,1 python test.py --img_dir '/opt/data/private/LV_Dataset/DiffGLV-Test-All/RealPhoto60/LQ' --save_dir ./results-F --SUPIR_sign F --upscale 2 --s_cfg 4.0 --linear_CFG
110
+ ```
111
+
112
+ ### Gradio Demo
113
+ ```Shell
114
+ CUDA_VISIBLE_DEVICES=0,1 python gradio_demo.py --ip 0.0.0.0 --port 6688 --use_image_slider --log_history
115
+
116
+ # Juggernaut_RunDiffusionPhoto2_Lightning_4Steps and DPM++ M2 SDE Karras for fast sampling
117
+ CUDA_VISIBLE_DEVICES=0,1 python gradio_demo.py --ip 0.0.0.0 --port 6688 --use_image_slider --log_history --opt options/SUPIR_v0_Juggernautv9_lightning.yaml
118
+
119
+ # less VRAM & slower (12G for Diffusion, 16G for LLaVA)
120
+ CUDA_VISIBLE_DEVICES=0,1 python gradio_demo.py --ip 0.0.0.0 --port 6688 --use_image_slider --log_history --loading_half_params --use_tile_vae --load_8bit_llava
121
+ ```
122
+ <p align="center">
123
+ <img src="assets/DemoGuide.png">
124
+ </p>
125
+
126
+
127
+ ### Online Demo (Coming Soon)
128
+
129
+
130
+ ---
131
+
132
+ ## BibTeX
133
+ @misc{yu2024scaling,
134
+ title={Scaling Up to Excellence: Practicing Model Scaling for Photo-Realistic Image Restoration In the Wild},
135
+ author={Fanghua Yu and Jinjin Gu and Zheyuan Li and Jinfan Hu and Xiangtao Kong and Xintao Wang and Jingwen He and Yu Qiao and Chao Dong},
136
+ year={2024},
137
+ eprint={2401.13627},
138
+ archivePrefix={arXiv},
139
+ primaryClass={cs.CV}
140
+ }
141
+
142
+ ---
143
+
144
+ ## 📧 Contact
145
+ If you have any question, please email `fanghuayu96@gmail.com`.
146
+
147
+ ---
148
+ ## Non-Commercial Use Only Declaration
149
+ The SUPIR ("Software") is made available for use, reproduction, and distribution strictly for non-commercial purposes. For the purposes of this declaration, "non-commercial" is defined as not primarily intended for or directed towards commercial advantage or monetary compensation.
150
+
151
+ By using, reproducing, or distributing the Software, you agree to abide by this restriction and not to use the Software for any commercial purposes without obtaining prior written permission from Dr. Jinjin Gu.
152
+
153
+ This declaration does not in any way limit the rights under any open source license that may apply to the Software; it solely adds a condition that the Software shall not be used for commercial purposes.
154
+
155
+ IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
156
+
157
+ For inquiries or to obtain permission for commercial use, please contact Dr. Jinjin Gu (hellojasongt@gmail.com).
gradio_demo.py ADDED
@@ -0,0 +1,314 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ import gradio as gr
4
+ from gradio_imageslider import ImageSlider
5
+ import argparse
6
+ from SUPIR.util import HWC3, upscale_image, fix_resize, convert_dtype
7
+ import numpy as np
8
+ import torch
9
+ from SUPIR.util import create_SUPIR_model, load_QF_ckpt
10
+ from PIL import Image
11
+ from llava.llava_agent import LLavaAgent
12
+ from CKPT_PTH import LLAVA_MODEL_PATH
13
+ import einops
14
+ import copy
15
+ import time
16
+
17
+ parser = argparse.ArgumentParser()
18
+ parser.add_argument("--opt", type=str, default='options/SUPIR_v0.yaml')
19
+ parser.add_argument("--ip", type=str, default='127.0.0.1')
20
+ parser.add_argument("--port", type=int, default='6688')
21
+ parser.add_argument("--no_llava", action='store_true', default=False)
22
+ parser.add_argument("--use_image_slider", action='store_true', default=False)
23
+ parser.add_argument("--log_history", action='store_true', default=False)
24
+ parser.add_argument("--loading_half_params", action='store_true', default=False)
25
+ parser.add_argument("--use_tile_vae", action='store_true', default=False)
26
+ parser.add_argument("--encoder_tile_size", type=int, default=512)
27
+ parser.add_argument("--decoder_tile_size", type=int, default=64)
28
+ parser.add_argument("--load_8bit_llava", action='store_true', default=False)
29
+ args = parser.parse_args()
30
+ server_ip = args.ip
31
+ server_port = args.port
32
+ use_llava = not args.no_llava
33
+
34
+ if torch.cuda.device_count() >= 2:
35
+ SUPIR_device = 'cuda:0'
36
+ LLaVA_device = 'cuda:1'
37
+ elif torch.cuda.device_count() == 1:
38
+ SUPIR_device = 'cuda:0'
39
+ LLaVA_device = 'cuda:0'
40
+ else:
41
+ raise ValueError('Currently support CUDA only.')
42
+
43
+ # load SUPIR
44
+ model, default_setting = create_SUPIR_model(args.opt, SUPIR_sign='Q', load_default_setting=True)
45
+ if args.loading_half_params:
46
+ model = model.half()
47
+ if args.use_tile_vae:
48
+ model.init_tile_vae(encoder_tile_size=args.encoder_tile_size, decoder_tile_size=args.decoder_tile_size)
49
+ model = model.to(SUPIR_device)
50
+ model.first_stage_model.denoise_encoder_s1 = copy.deepcopy(model.first_stage_model.denoise_encoder)
51
+ model.current_model = 'v0-Q'
52
+ ckpt_Q, ckpt_F = load_QF_ckpt(args.opt)
53
+
54
+ # load LLaVA
55
+ if use_llava:
56
+ llava_agent = LLavaAgent(LLAVA_MODEL_PATH, device=LLaVA_device, load_8bit=args.load_8bit_llava, load_4bit=False)
57
+ else:
58
+ llava_agent = None
59
+
60
+ def stage1_process(input_image, gamma_correction):
61
+ torch.cuda.set_device(SUPIR_device)
62
+ LQ = HWC3(input_image)
63
+ LQ = fix_resize(LQ, 512)
64
+ # stage1
65
+ LQ = np.array(LQ) / 255 * 2 - 1
66
+ LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
67
+ LQ = model.batchify_denoise(LQ, is_stage1=True)
68
+ LQ = (LQ[0].permute(1, 2, 0) * 127.5 + 127.5).cpu().numpy().round().clip(0, 255).astype(np.uint8)
69
+ # gamma correction
70
+ LQ = LQ / 255.0
71
+ LQ = np.power(LQ, gamma_correction)
72
+ LQ *= 255.0
73
+ LQ = LQ.round().clip(0, 255).astype(np.uint8)
74
+ return LQ
75
+
76
+ def llave_process(input_image, temperature, top_p, qs=None):
77
+ torch.cuda.set_device(LLaVA_device)
78
+ if use_llava:
79
+ LQ = HWC3(input_image)
80
+ LQ = Image.fromarray(LQ.astype('uint8'))
81
+ captions = llava_agent.gen_image_caption([LQ], temperature=temperature, top_p=top_p, qs=qs)
82
+ else:
83
+ captions = ['LLaVA is not available. Please add text manually.']
84
+ return captions[0]
85
+
86
+ def stage2_process(input_image, prompt, a_prompt, n_prompt, num_samples, upscale, edm_steps, s_stage1, s_stage2,
87
+ s_cfg, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction,
88
+ linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select):
89
+ torch.cuda.set_device(SUPIR_device)
90
+ event_id = str(time.time_ns())
91
+ event_dict = {'event_id': event_id, 'localtime': time.ctime(), 'prompt': prompt, 'a_prompt': a_prompt,
92
+ 'n_prompt': n_prompt, 'num_samples': num_samples, 'upscale': upscale, 'edm_steps': edm_steps,
93
+ 's_stage1': s_stage1, 's_stage2': s_stage2, 's_cfg': s_cfg, 'seed': seed, 's_churn': s_churn,
94
+ 's_noise': s_noise, 'color_fix_type': color_fix_type, 'diff_dtype': diff_dtype, 'ae_dtype': ae_dtype,
95
+ 'gamma_correction': gamma_correction, 'linear_CFG': linear_CFG, 'linear_s_stage2': linear_s_stage2,
96
+ 'spt_linear_CFG': spt_linear_CFG, 'spt_linear_s_stage2': spt_linear_s_stage2,
97
+ 'model_select': model_select}
98
+
99
+ if model_select != model.current_model:
100
+ if model_select == 'v0-Q':
101
+ print('load v0-Q')
102
+ model.load_state_dict(ckpt_Q, strict=False)
103
+ model.current_model = 'v0-Q'
104
+ elif model_select == 'v0-F':
105
+ print('load v0-F')
106
+ model.load_state_dict(ckpt_F, strict=False)
107
+ model.current_model = 'v0-F'
108
+ input_image = HWC3(input_image)
109
+ input_image = upscale_image(input_image, upscale, unit_resolution=32,
110
+ min_size=1024)
111
+
112
+ LQ = np.array(input_image) / 255.0
113
+ LQ = np.power(LQ, gamma_correction)
114
+ LQ *= 255.0
115
+ LQ = LQ.round().clip(0, 255).astype(np.uint8)
116
+ LQ = LQ / 255 * 2 - 1
117
+ LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
118
+ if use_llava:
119
+ captions = [prompt]
120
+ else:
121
+ captions = ['']
122
+
123
+ model.ae_dtype = convert_dtype(ae_dtype)
124
+ model.model.dtype = convert_dtype(diff_dtype)
125
+
126
+ samples = model.batchify_sample(LQ, captions, num_steps=edm_steps, restoration_scale=s_stage1, s_churn=s_churn,
127
+ s_noise=s_noise, cfg_scale=s_cfg, control_scale=s_stage2, seed=seed,
128
+ num_samples=num_samples, p_p=a_prompt, n_p=n_prompt, color_fix_type=color_fix_type,
129
+ use_linear_CFG=linear_CFG, use_linear_control_scale=linear_s_stage2,
130
+ cfg_scale_start=spt_linear_CFG, control_scale_start=spt_linear_s_stage2)
131
+
132
+ x_samples = (einops.rearrange(samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().round().clip(
133
+ 0, 255).astype(np.uint8)
134
+ results = [x_samples[i] for i in range(num_samples)]
135
+
136
+ if args.log_history:
137
+ os.makedirs(f'./history/{event_id[:5]}/{event_id[5:]}', exist_ok=True)
138
+ with open(f'./history/{event_id[:5]}/{event_id[5:]}/logs.txt', 'w') as f:
139
+ f.write(str(event_dict))
140
+ f.close()
141
+ Image.fromarray(input_image).save(f'./history/{event_id[:5]}/{event_id[5:]}/LQ.png')
142
+ for i, result in enumerate(results):
143
+ Image.fromarray(result).save(f'./history/{event_id[:5]}/{event_id[5:]}/HQ_{i}.png')
144
+ return [input_image] + results, event_id, 3, ''
145
+
146
+
147
+ def load_and_reset(param_setting):
148
+ edm_steps = default_setting.edm_steps
149
+ s_stage2 = 1.0
150
+ s_stage1 = -1.0
151
+ s_churn = 5
152
+ s_noise = 1.003
153
+ a_prompt = 'Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - ' \
154
+ 'realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore ' \
155
+ 'detailing, hyper sharpness, perfect without deformations.'
156
+ n_prompt = 'painting, oil painting, illustration, drawing, art, sketch, oil painting, cartoon, CG Style, ' \
157
+ '3D render, unreal engine, blurring, dirty, messy, worst quality, low quality, frames, watermark, ' \
158
+ 'signature, jpeg artifacts, deformed, lowres, over-smooth'
159
+ color_fix_type = 'Wavelet'
160
+ spt_linear_s_stage2 = 0.0
161
+ linear_s_stage2 = False
162
+ linear_CFG = True
163
+ if param_setting == "Quality":
164
+ s_cfg = default_setting.s_cfg_Quality
165
+ spt_linear_CFG = default_setting.spt_linear_CFG_Quality
166
+ elif param_setting == "Fidelity":
167
+ s_cfg = default_setting.s_cfg_Fidelity
168
+ spt_linear_CFG = default_setting.spt_linear_CFG_Fidelity
169
+ else:
170
+ raise NotImplementedError
171
+ return edm_steps, s_cfg, s_stage2, s_stage1, s_churn, s_noise, a_prompt, n_prompt, color_fix_type, linear_CFG, \
172
+ linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2
173
+
174
+
175
+ def submit_feedback(event_id, fb_score, fb_text):
176
+ if args.log_history:
177
+ with open(f'./history/{event_id[:5]}/{event_id[5:]}/logs.txt', 'r') as f:
178
+ event_dict = eval(f.read())
179
+ f.close()
180
+ event_dict['feedback'] = {'score': fb_score, 'text': fb_text}
181
+ with open(f'./history/{event_id[:5]}/{event_id[5:]}/logs.txt', 'w') as f:
182
+ f.write(str(event_dict))
183
+ f.close()
184
+ return 'Submit successfully, thank you for your comments!'
185
+ else:
186
+ return 'Submit failed, the server is not set to log history.'
187
+
188
+
189
+ title_md = """
190
+ # **SUPIR: Practicing Model Scaling for Photo-Realistic Image Restoration**
191
+
192
+ ⚠️SUPIR is still a research project under tested and is not yet a stable commercial product.
193
+
194
+ [[Paper](https://arxiv.org/abs/2401.13627)] &emsp; [[Project Page](http://supir.xpixel.group/)] &emsp; [[How to play](https://github.com/Fanghua-Yu/SUPIR/blob/master/assets/DemoGuide.png)]
195
+ """
196
+
197
+
198
+ claim_md = """
199
+ ## **Terms of use**
200
+
201
+ By using this service, users are required to agree to the following terms: The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research. Please submit a feedback to us if you get any inappropriate answer! We will collect those to keep improving our models. For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality.
202
+
203
+ ## **License**
204
+
205
+ The service is a research preview intended for non-commercial use only, subject to the model [License](https://github.com/Fanghua-Yu/SUPIR) of SUPIR.
206
+ """
207
+
208
+
209
+ block = gr.Blocks(title='SUPIR').queue()
210
+ with block:
211
+ with gr.Row():
212
+ gr.Markdown(title_md)
213
+ with gr.Row():
214
+ with gr.Column():
215
+ with gr.Row(equal_height=True):
216
+ with gr.Column():
217
+ gr.Markdown("<center>Input</center>")
218
+ input_image = gr.Image(type="numpy", elem_id="image-input", height=400, width=400)
219
+ with gr.Column():
220
+ gr.Markdown("<center>Stage1 Output</center>")
221
+ denoise_image = gr.Image(type="numpy", elem_id="image-s1", height=400, width=400)
222
+ prompt = gr.Textbox(label="Prompt", value="")
223
+ with gr.Accordion("Stage1 options", open=False):
224
+ gamma_correction = gr.Slider(label="Gamma Correction", minimum=0.1, maximum=2.0, value=1.0, step=0.1)
225
+ with gr.Accordion("LLaVA options", open=False):
226
+ temperature = gr.Slider(label="Temperature", minimum=0., maximum=1.0, value=0.2, step=0.1)
227
+ top_p = gr.Slider(label="Top P", minimum=0., maximum=1.0, value=0.7, step=0.1)
228
+ qs = gr.Textbox(label="Question", value="Describe this image and its style in a very detailed manner. "
229
+ "The image is a realistic photography, not an art painting.")
230
+ with gr.Accordion("Stage2 options", open=False):
231
+ num_samples = gr.Slider(label="Num Samples", minimum=1, maximum=4 if not args.use_image_slider else 1
232
+ , value=1, step=1)
233
+ upscale = gr.Slider(label="Upscale", minimum=1, maximum=8, value=1, step=1)
234
+ edm_steps = gr.Slider(label="Steps", minimum=1, maximum=200, value=default_setting.edm_steps, step=1)
235
+ s_cfg = gr.Slider(label="Text Guidance Scale", minimum=1.0, maximum=15.0,
236
+ value=default_setting.s_cfg_Quality, step=0.1)
237
+ s_stage2 = gr.Slider(label="Stage2 Guidance Strength", minimum=0., maximum=1., value=1., step=0.05)
238
+ s_stage1 = gr.Slider(label="Stage1 Guidance Strength", minimum=-1.0, maximum=6.0, value=-1.0, step=1.0)
239
+ seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
240
+ s_churn = gr.Slider(label="S-Churn", minimum=0, maximum=40, value=5, step=1)
241
+ s_noise = gr.Slider(label="S-Noise", minimum=1.0, maximum=1.1, value=1.003, step=0.001)
242
+ a_prompt = gr.Textbox(label="Default Positive Prompt",
243
+ value='Cinematic, High Contrast, highly detailed, taken using a Canon EOS R '
244
+ 'camera, hyper detailed photo - realistic maximum detail, 32k, Color '
245
+ 'Grading, ultra HD, extreme meticulous detailing, skin pore detailing, '
246
+ 'hyper sharpness, perfect without deformations.')
247
+ n_prompt = gr.Textbox(label="Default Negative Prompt",
248
+ value='painting, oil painting, illustration, drawing, art, sketch, oil painting, '
249
+ 'cartoon, CG Style, 3D render, unreal engine, blurring, dirty, messy, '
250
+ 'worst quality, low quality, frames, watermark, signature, jpeg artifacts, '
251
+ 'deformed, lowres, over-smooth')
252
+ with gr.Row():
253
+ with gr.Column():
254
+ linear_CFG = gr.Checkbox(label="Linear CFG", value=True)
255
+ spt_linear_CFG = gr.Slider(label="CFG Start", minimum=1.0,
256
+ maximum=9.0, value=default_setting.spt_linear_CFG_Quality, step=0.5)
257
+ with gr.Column():
258
+ linear_s_stage2 = gr.Checkbox(label="Linear Stage2 Guidance", value=False)
259
+ spt_linear_s_stage2 = gr.Slider(label="Guidance Start", minimum=0.,
260
+ maximum=1., value=0., step=0.05)
261
+ with gr.Row():
262
+ with gr.Column():
263
+ diff_dtype = gr.Radio(['fp32', 'fp16', 'bf16'], label="Diffusion Data Type", value="fp16",
264
+ interactive=True)
265
+ with gr.Column():
266
+ ae_dtype = gr.Radio(['fp32', 'bf16'], label="Auto-Encoder Data Type", value="bf16",
267
+ interactive=True)
268
+ with gr.Column():
269
+ color_fix_type = gr.Radio(["None", "AdaIn", "Wavelet"], label="Color-Fix Type", value="Wavelet",
270
+ interactive=True)
271
+ with gr.Column():
272
+ model_select = gr.Radio(["v0-Q", "v0-F"], label="Model Selection", value="v0-Q",
273
+ interactive=True)
274
+
275
+ with gr.Column():
276
+ gr.Markdown("<center>Stage2 Output</center>")
277
+ if not args.use_image_slider:
278
+ result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery1")
279
+ else:
280
+ result_gallery = ImageSlider(label='Output', show_label=False, elem_id="gallery1")
281
+ with gr.Row():
282
+ with gr.Column():
283
+ denoise_button = gr.Button(value="Stage1 Run")
284
+ with gr.Column():
285
+ llave_button = gr.Button(value="LlaVa Run")
286
+ with gr.Column():
287
+ diffusion_button = gr.Button(value="Stage2 Run")
288
+ with gr.Row():
289
+ with gr.Column():
290
+ param_setting = gr.Dropdown(["Quality", "Fidelity"], interactive=True, label="Param Setting",
291
+ value="Quality")
292
+ with gr.Column():
293
+ restart_button = gr.Button(value="Reset Param", scale=2)
294
+ with gr.Accordion("Feedback", open=True):
295
+ fb_score = gr.Slider(label="Feedback Score", minimum=1, maximum=5, value=3, step=1,
296
+ interactive=True)
297
+ fb_text = gr.Textbox(label="Feedback Text", value="", placeholder='Please enter your feedback here.')
298
+ submit_button = gr.Button(value="Submit Feedback")
299
+ with gr.Row():
300
+ gr.Markdown(claim_md)
301
+ event_id = gr.Textbox(label="Event ID", value="", visible=False)
302
+
303
+ llave_button.click(fn=llave_process, inputs=[denoise_image, temperature, top_p, qs], outputs=[prompt])
304
+ denoise_button.click(fn=stage1_process, inputs=[input_image, gamma_correction],
305
+ outputs=[denoise_image])
306
+ stage2_ips = [input_image, prompt, a_prompt, n_prompt, num_samples, upscale, edm_steps, s_stage1, s_stage2,
307
+ s_cfg, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction,
308
+ linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select]
309
+ diffusion_button.click(fn=stage2_process, inputs=stage2_ips, outputs=[result_gallery, event_id, fb_score, fb_text])
310
+ restart_button.click(fn=load_and_reset, inputs=[param_setting],
311
+ outputs=[edm_steps, s_cfg, s_stage2, s_stage1, s_churn, s_noise, a_prompt, n_prompt,
312
+ color_fix_type, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2])
313
+ submit_button.click(fn=submit_feedback, inputs=[event_id, fb_score, fb_text], outputs=[fb_text])
314
+ block.launch(server_name=server_ip, server_port=server_port)
gradio_demo_face.py ADDED
@@ -0,0 +1,411 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ import gradio as gr
4
+ from gradio_imageslider import ImageSlider
5
+ import argparse
6
+ from SUPIR.util import HWC3, upscale_image, fix_resize, convert_dtype
7
+ import numpy as np
8
+ import torch
9
+ from SUPIR.util import create_SUPIR_model, load_QF_ckpt
10
+ from PIL import Image
11
+ from llava.llava_agent import LLavaAgent
12
+ from CKPT_PTH import LLAVA_MODEL_PATH
13
+ import einops
14
+ import copy
15
+ import time
16
+ from omegaconf import OmegaConf
17
+ from SUPIR.utils.face_restoration_helper import FaceRestoreHelper
18
+
19
+ parser = argparse.ArgumentParser()
20
+ parser.add_argument("--ip", type=str, default='127.0.0.1')
21
+ parser.add_argument("--port", type=int, default='6688')
22
+ parser.add_argument("--no_llava", action='store_true', default=False)
23
+ parser.add_argument("--use_image_slider", action='store_true', default=False)
24
+ parser.add_argument("--log_history", action='store_true', default=False)
25
+ parser.add_argument("--loading_half_params", action='store_true', default=False)
26
+ parser.add_argument("--use_tile_vae", action='store_true', default=False)
27
+ parser.add_argument("--load_8bit_llava", action='store_true', default=False)
28
+ parser.add_argument("--local_prompt", action='store_true', default=False)
29
+ args = parser.parse_args()
30
+ server_ip = args.ip
31
+ server_port = args.port
32
+ use_llava = not args.no_llava
33
+
34
+ if torch.cuda.device_count() >= 2:
35
+ SUPIR_device = 'cuda:0'
36
+ LLaVA_device = 'cuda:1'
37
+ elif torch.cuda.device_count() == 1:
38
+ SUPIR_device = 'cuda:0'
39
+ LLaVA_device = 'cuda:0'
40
+ else:
41
+ raise ValueError('Currently support CUDA only.')
42
+
43
+ # load SUPIR
44
+ config_path = 'options/SUPIR_v0.yaml'
45
+ config = OmegaConf.load(config_path)
46
+ model = create_SUPIR_model(config_path, SUPIR_sign='Q')
47
+ if args.loading_half_params:
48
+ model = model.half()
49
+ if args.use_tile_vae:
50
+ model.init_tile_vae(encoder_tile_size=512, decoder_tile_size=64)
51
+ model = model.to(SUPIR_device)
52
+ model.first_stage_model.denoise_encoder_s1 = copy.deepcopy(model.first_stage_model.denoise_encoder)
53
+ model.current_model = 'v0-Q'
54
+ ckpt_Q, ckpt_F = load_QF_ckpt('options/SUPIR_v0.yaml')
55
+
56
+ # load LLaVA
57
+ if use_llava:
58
+ llava_agent = LLavaAgent(LLAVA_MODEL_PATH, device=LLaVA_device, load_8bit=args.load_8bit_llava, load_4bit=False)
59
+ else:
60
+ llava_agent = None
61
+
62
+ # load face helper
63
+ face_helper = FaceRestoreHelper(
64
+ device=SUPIR_device,
65
+ upscale_factor=1,
66
+ face_size=1024,
67
+ use_parse=True,
68
+ det_model='retinaface_resnet50'
69
+ )
70
+
71
+ # only exhibit the overall quality of the stage1 output
72
+ def stage1_process(input_image, gamma_correction):
73
+ torch.cuda.set_device(SUPIR_device)
74
+ LQ = HWC3(input_image)
75
+ LQ = fix_resize(LQ, 512)
76
+ # stage1
77
+ LQ = np.array(LQ) / 255 * 2 - 1
78
+ LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
79
+ LQ = model.batchify_denoise(LQ, is_stage1=True)
80
+ LQ = (LQ[0].permute(1, 2, 0) * 127.5 + 127.5).cpu().numpy().round().clip(0, 255).astype(np.uint8)
81
+ # gamma correction
82
+ LQ = LQ / 255.0
83
+ LQ = np.power(LQ, gamma_correction)
84
+ LQ *= 255.0
85
+ LQ = LQ.round().clip(0, 255).astype(np.uint8)
86
+ return LQ
87
+
88
+ def llave_process(input_image, upscale, temperature, top_p, qs=None):
89
+ torch.cuda.set_device(SUPIR_device)
90
+ input_image = HWC3(input_image)
91
+ input_image = upscale_image(input_image, upscale, unit_resolution=32,
92
+ min_size=1024)
93
+ LQ = np.array(input_image) / 255 * 2 - 1
94
+ LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
95
+ LQ = model.batchify_denoise(LQ, is_stage1=True)
96
+
97
+ LQ = (LQ[0].permute(1, 2, 0) * 127.5 + 127.5).cpu().numpy().round().clip(0, 255).astype(np.uint8)
98
+ LQs = [Image.fromarray(LQ)]
99
+
100
+ face_helper.clean_all()
101
+ face_helper.read_image(LQ)
102
+ # get face landmarks for each face
103
+ face_helper.get_face_landmarks_5(only_center_face=False, resize=640, eye_dist_threshold=5)
104
+ face_helper.align_warp_face()
105
+
106
+ for face in face_helper.cropped_faces:
107
+ LQs.append(Image.fromarray(face))
108
+
109
+ captions = []
110
+ torch.cuda.set_device(LLaVA_device)
111
+ if use_llava:
112
+ for LQ in LQs:
113
+ captions += llava_agent.gen_image_caption([LQ], temperature=temperature, top_p=top_p, qs=qs)
114
+ else:
115
+ captions = ['LLaVA is not available. Please add text manually.']
116
+ del LQs[0]
117
+ return str(captions), [np.array(face) for face in LQs]
118
+
119
+
120
+ def stage2_process(input_image, prompt, a_prompt, n_prompt, num_samples, upscale, edm_steps, s_stage1, s_stage2,
121
+ s_cfg, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction,
122
+ linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select,
123
+ face_resolution, apply_bg, apply_face):
124
+ torch.cuda.set_device(SUPIR_device)
125
+ event_id = str(time.time_ns())
126
+ event_dict = {'event_id': event_id, 'localtime': time.ctime(), 'prompt': prompt, 'a_prompt': a_prompt,
127
+ 'n_prompt': n_prompt, 'num_samples': num_samples, 'upscale': upscale, 'edm_steps': edm_steps,
128
+ 's_stage1': s_stage1, 's_stage2': s_stage2, 's_cfg': s_cfg, 'seed': seed, 's_churn': s_churn,
129
+ 's_noise': s_noise, 'color_fix_type': color_fix_type, 'diff_dtype': diff_dtype, 'ae_dtype': ae_dtype,
130
+ 'gamma_correction': gamma_correction, 'linear_CFG': linear_CFG, 'linear_s_stage2': linear_s_stage2,
131
+ 'spt_linear_CFG': spt_linear_CFG, 'spt_linear_s_stage2': spt_linear_s_stage2,
132
+ 'model_select': model_select}
133
+
134
+ if model_select != model.current_model:
135
+ if model_select == 'v0-Q':
136
+ print('load v0-Q')
137
+ model.load_state_dict(ckpt_Q, strict=False)
138
+ model.current_model = 'v0-Q'
139
+ elif model_select == 'v0-F':
140
+ print('load v0-F')
141
+ model.load_state_dict(ckpt_F, strict=False)
142
+ model.current_model = 'v0-F'
143
+ input_image = HWC3(input_image)
144
+ input_image = upscale_image(input_image, upscale, unit_resolution=32,
145
+ min_size=1024)
146
+
147
+ LQ = np.array(input_image)
148
+ face_helper.clean_all()
149
+ face_helper.read_image(LQ)
150
+ # get face landmarks for each face
151
+ face_helper.get_face_landmarks_5(only_center_face=False, resize=640, eye_dist_threshold=5)
152
+ face_helper.align_warp_face()
153
+
154
+ LQ = LQ / 255 * 2 - 1
155
+ LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
156
+
157
+ if use_llava and prompt != '':
158
+ captions = eval(prompt)
159
+ else:
160
+ captions = [''] * (1 + len(face_helper.cropped_faces))
161
+
162
+ bg_caption, face_captions = captions[0], captions[1:]
163
+
164
+ model.ae_dtype = convert_dtype(ae_dtype)
165
+ model.model.dtype = convert_dtype(diff_dtype)
166
+
167
+ _faces = []
168
+ if apply_face:
169
+ faces = []
170
+ for face in face_helper.cropped_faces:
171
+ _faces.append(face)
172
+ face = np.array(face) / 255 * 2 - 1
173
+ face = torch.tensor(face, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
174
+ faces.append(face)
175
+
176
+ for face, caption in zip(faces, face_captions):
177
+ caption = [caption]
178
+
179
+ from torch.nn.functional import interpolate
180
+ face = interpolate(face, size=face_resolution, mode='bilinear', align_corners=False)
181
+ if face_resolution < 1024:
182
+ face = torch.nn.functional.pad(face, (512-face_resolution//2, 512-face_resolution//2,
183
+ 512-face_resolution//2, 512-face_resolution//2), 'constant', 0)
184
+
185
+ samples = model.batchify_sample(face, caption, num_steps=edm_steps, restoration_scale=s_stage1, s_churn=s_churn,
186
+ s_noise=s_noise, cfg_scale=s_cfg, control_scale=s_stage2, seed=seed,
187
+ num_samples=num_samples, p_p=a_prompt, n_p=n_prompt, color_fix_type=color_fix_type,
188
+ use_linear_CFG=linear_CFG, use_linear_control_scale=linear_s_stage2,
189
+ cfg_scale_start=spt_linear_CFG, control_scale_start=spt_linear_s_stage2)
190
+ if face_resolution < 1024:
191
+ samples = samples[:, :, 512-face_resolution//2:512+face_resolution//2,
192
+ 512-face_resolution//2:512+face_resolution//2]
193
+ samples = interpolate(samples, size=face_helper.face_size, mode='bilinear', align_corners=False)
194
+ x_samples = (einops.rearrange(samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().round().clip(
195
+ 0, 255).astype(np.uint8)
196
+
197
+ face_helper.add_restored_face(x_samples[0])
198
+ _faces.append(x_samples[0])
199
+
200
+ if apply_bg:
201
+ caption = [bg_caption]
202
+ samples = model.batchify_sample(LQ, caption, num_steps=edm_steps, restoration_scale=s_stage1,
203
+ s_churn=s_churn,
204
+ s_noise=s_noise, cfg_scale=s_cfg, control_scale=s_stage2, seed=seed,
205
+ num_samples=num_samples, p_p=a_prompt, n_p=n_prompt,
206
+ color_fix_type=color_fix_type,
207
+ use_linear_CFG=linear_CFG, use_linear_control_scale=linear_s_stage2,
208
+ cfg_scale_start=spt_linear_CFG, control_scale_start=spt_linear_s_stage2)
209
+ else:
210
+ samples = LQ
211
+ _bg = (einops.rearrange(samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().round().clip(
212
+ 0, 255).astype(np.uint8)
213
+ face_helper.get_inverse_affine(None)
214
+ results = [face_helper.paste_faces_to_input_image(upsample_img=_bg[0])]
215
+ else:
216
+ caption = [bg_caption]
217
+ samples = model.batchify_sample(LQ, caption, num_steps=edm_steps, restoration_scale=s_stage1,
218
+ s_churn=s_churn,
219
+ s_noise=s_noise, cfg_scale=s_cfg, control_scale=s_stage2, seed=seed,
220
+ num_samples=num_samples, p_p=a_prompt, n_p=n_prompt,
221
+ color_fix_type=color_fix_type,
222
+ use_linear_CFG=linear_CFG, use_linear_control_scale=linear_s_stage2,
223
+ cfg_scale_start=spt_linear_CFG, control_scale_start=spt_linear_s_stage2)
224
+ x_samples = (einops.rearrange(samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().round().clip(
225
+ 0, 255).astype(np.uint8)
226
+ results = [x_samples[0]]
227
+
228
+ if args.log_history:
229
+ os.makedirs(f'./history/{event_id[:5]}/{event_id[5:]}', exist_ok=True)
230
+ with open(f'./history/{event_id[:5]}/{event_id[5:]}/logs.txt', 'w') as f:
231
+ f.write(str(event_dict))
232
+ f.close()
233
+ Image.fromarray(input_image).save(f'./history/{event_id[:5]}/{event_id[5:]}/LQ.png')
234
+ for i, result in enumerate(results):
235
+ Image.fromarray(result).save(f'./history/{event_id[:5]}/{event_id[5:]}/HQ_{i}.png')
236
+ return [input_image] + results, event_id, 3, '', _faces
237
+
238
+ def load_and_reset(param_setting):
239
+ edm_steps = 50
240
+ s_stage2 = 1.0
241
+ s_stage1 = -1.0
242
+ s_churn = 5
243
+ s_noise = 1.003
244
+ a_prompt = 'Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - ' \
245
+ 'realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore ' \
246
+ 'detailing, hyper sharpness, perfect without deformations.'
247
+ n_prompt = 'painting, oil painting, illustration, drawing, art, sketch, oil painting, cartoon, CG Style, ' \
248
+ '3D render, unreal engine, blurring, dirty, messy, worst quality, low quality, frames, watermark, ' \
249
+ 'signature, jpeg artifacts, deformed, lowres, over-smooth'
250
+ color_fix_type = 'Wavelet'
251
+ spt_linear_s_stage2 = 0.0
252
+ linear_s_stage2 = False
253
+ linear_CFG = True
254
+ if param_setting == "Quality":
255
+ s_cfg = 7.5
256
+ spt_linear_CFG = 4.0
257
+ elif param_setting == "Fidelity":
258
+ s_cfg = 4.0
259
+ spt_linear_CFG = 1.0
260
+ else:
261
+ raise NotImplementedError
262
+ return edm_steps, s_cfg, s_stage2, s_stage1, s_churn, s_noise, a_prompt, n_prompt, color_fix_type, linear_CFG, \
263
+ linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2
264
+
265
+
266
+ def submit_feedback(event_id, fb_score, fb_text):
267
+ if args.log_history:
268
+ with open(f'./history/{event_id[:5]}/{event_id[5:]}/logs.txt', 'r') as f:
269
+ event_dict = eval(f.read())
270
+ f.close()
271
+ event_dict['feedback'] = {'score': fb_score, 'text': fb_text}
272
+ with open(f'./history/{event_id[:5]}/{event_id[5:]}/logs.txt', 'w') as f:
273
+ f.write(str(event_dict))
274
+ f.close()
275
+ return 'Submit successfully, thank you for your comments!'
276
+ else:
277
+ return 'Submit failed, the server is not set to log history.'
278
+
279
+ title_md = """
280
+ # **SUPIR: Practicing Model Scaling for Photo-Realistic Image Restoration**
281
+
282
+ ⚠️SUPIR is still a research project under tested and is not yet a stable commercial product.
283
+
284
+ [[Paper](https://arxiv.org/abs/2401.13627)] &emsp; [[Project Page](http://supir.xpixel.group/)] &emsp; [[How to play](https://github.com/Fanghua-Yu/SUPIR/blob/master/assets/DemoGuide.png)]
285
+ """
286
+
287
+
288
+ claim_md = """
289
+ ## **Terms of use**
290
+
291
+ By using this service, users are required to agree to the following terms: The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research. Please submit a feedback to us if you get any inappropriate answer! We will collect those to keep improving our models. For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality.
292
+
293
+ ## **License**
294
+
295
+ The service is a research preview intended for non-commercial use only, subject to the model [License](https://github.com/Fanghua-Yu/SUPIR) of SUPIR.
296
+ """
297
+
298
+
299
+ block = gr.Blocks(title='SUPIR').queue()
300
+ with block:
301
+ with gr.Row():
302
+ gr.Markdown(title_md)
303
+ with gr.Row():
304
+ with gr.Column():
305
+ with gr.Row(equal_height=True):
306
+ with gr.Column():
307
+ gr.Markdown("<center>Input</center>")
308
+ input_image = gr.Image(type="numpy", elem_id="image-input", height=400, width=400)
309
+ with gr.Column():
310
+ gr.Markdown("<center>Stage1 Output</center>")
311
+ denoise_image = gr.Image(type="numpy", elem_id="image-s1", height=400, width=400)
312
+ prompt = gr.Textbox(label="Prompt", value="")
313
+ with gr.Accordion("Stage1 options", open=False):
314
+ gamma_correction = gr.Slider(label="Gamma Correction", minimum=0.1, maximum=2.0, value=1.0, step=0.1)
315
+ with gr.Accordion("LLaVA options", open=False):
316
+ temperature = gr.Slider(label="Temperature", minimum=0., maximum=1.0, value=0.2, step=0.1)
317
+ top_p = gr.Slider(label="Top P", minimum=0., maximum=1.0, value=0.7, step=0.1)
318
+ qs = gr.Textbox(label="Question", value="Describe this image and its style in a very detailed manner. "
319
+ "The image is a realistic photography, not an art painting.")
320
+ with gr.Accordion("Stage2 options", open=False):
321
+ num_samples = gr.Slider(label="Num Samples", minimum=1, maximum=4 if not args.use_image_slider else 1
322
+ , value=1, step=1)
323
+ upscale = gr.Slider(label="Upscale", minimum=1, maximum=8, value=1, step=1)
324
+ edm_steps = gr.Slider(label="Steps", minimum=20, maximum=200, value=50, step=1)
325
+ s_cfg = gr.Slider(label="Text Guidance Scale", minimum=1.0, maximum=15.0, value=7.5, step=0.1)
326
+ s_stage2 = gr.Slider(label="Stage2 Guidance Strength", minimum=0., maximum=1., value=1., step=0.05)
327
+ s_stage1 = gr.Slider(label="Stage1 Guidance Strength", minimum=-1.0, maximum=6.0, value=-1.0, step=1.0)
328
+ seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
329
+ s_churn = gr.Slider(label="S-Churn", minimum=0, maximum=40, value=5, step=1)
330
+ s_noise = gr.Slider(label="S-Noise", minimum=1.0, maximum=1.1, value=1.003, step=0.001)
331
+ a_prompt = gr.Textbox(label="Default Positive Prompt",
332
+ value='Cinematic, High Contrast, highly detailed, taken using a Canon EOS R '
333
+ 'camera, hyper detailed photo - realistic maximum detail, 32k, Color '
334
+ 'Grading, ultra HD, extreme meticulous detailing, skin pore detailing, '
335
+ 'hyper sharpness, perfect without deformations.')
336
+ n_prompt = gr.Textbox(label="Default Negative Prompt",
337
+ value='painting, oil painting, illustration, drawing, art, sketch, oil painting, '
338
+ 'cartoon, CG Style, 3D render, unreal engine, blurring, dirty, messy, '
339
+ 'worst quality, low quality, frames, watermark, signature, jpeg artifacts, '
340
+ 'deformed, lowres, over-smooth')
341
+ with gr.Row():
342
+ with gr.Column():
343
+ linear_CFG = gr.Checkbox(label="Linear CFG", value=True)
344
+ spt_linear_CFG = gr.Slider(label="CFG Start", minimum=1.0,
345
+ maximum=9.0, value=4.0, step=0.5)
346
+ with gr.Column():
347
+ linear_s_stage2 = gr.Checkbox(label="Linear Stage2 Guidance", value=False)
348
+ spt_linear_s_stage2 = gr.Slider(label="Guidance Start", minimum=0.,
349
+ maximum=1., value=0., step=0.05)
350
+ with gr.Row():
351
+ with gr.Column():
352
+ diff_dtype = gr.Radio(['fp32', 'fp16', 'bf16'], label="Diffusion Data Type", value="fp16",
353
+ interactive=True)
354
+ with gr.Column():
355
+ ae_dtype = gr.Radio(['fp32', 'bf16'], label="Auto-Encoder Data Type", value="bf16",
356
+ interactive=True)
357
+ with gr.Column():
358
+ color_fix_type = gr.Radio(["None", "AdaIn", "Wavelet"], label="Color-Fix Type", value="Wavelet",
359
+ interactive=True)
360
+ with gr.Column():
361
+ model_select = gr.Radio(["v0-Q", "v0-F"], label="Model Selection", value="v0-Q",
362
+ interactive=True)
363
+
364
+ with gr.Column():
365
+ gr.Markdown("<center>Stage2 Output</center>")
366
+ if not args.use_image_slider:
367
+ result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery1")
368
+ else:
369
+ result_gallery = ImageSlider(label='Output', show_label=False, elem_id="gallery1")
370
+ with gr.Row():
371
+ with gr.Column():
372
+ denoise_button = gr.Button(value="Stage1 Run")
373
+ with gr.Column():
374
+ llave_button = gr.Button(value="LlaVa Run")
375
+ with gr.Column():
376
+ diffusion_button = gr.Button(value="Stage2 Run")
377
+ with gr.Row():
378
+ with gr.Column():
379
+ param_setting = gr.Dropdown(["Quality", "Fidelity"], interactive=True, label="Param Setting",
380
+ value="Quality")
381
+ with gr.Column():
382
+ restart_button = gr.Button(value="Reset Param", scale=2)
383
+ with gr.Accordion("Face Options", open=True):
384
+ face_resolution = gr.Slider(label="Text Guidance Scale", minimum=256, maximum=2048, value=1024, step=32)
385
+ with gr.Row():
386
+ with gr.Column():
387
+ apply_bg = gr.Checkbox(label="BG restoration", value=True)
388
+ with gr.Column():
389
+ apply_face = gr.Checkbox(label="Face restoration", value=True)
390
+ with gr.Accordion("Feedback", open=False):
391
+ fb_score = gr.Slider(label="Feedback Score", minimum=1, maximum=5, value=3, step=1,
392
+ interactive=True)
393
+ fb_text = gr.Textbox(label="Feedback Text", value="", placeholder='Please enter your feedback here.')
394
+ submit_button = gr.Button(value="Submit Feedback")
395
+ face_gallery = gr.Gallery(label='Faces', show_label=False, elem_id="gallery2")
396
+ with gr.Row():
397
+ gr.Markdown(claim_md)
398
+ event_id = gr.Textbox(label="Event ID", value="", visible=False)
399
+
400
+ llave_button.click(fn=llave_process, inputs=[input_image, upscale, temperature, top_p, qs], outputs=[prompt, face_gallery])
401
+ denoise_button.click(fn=stage1_process, inputs=[input_image, gamma_correction],
402
+ outputs=[denoise_image])
403
+ stage2_ips = [input_image, prompt, a_prompt, n_prompt, num_samples, upscale, edm_steps, s_stage1, s_stage2,
404
+ s_cfg, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction,
405
+ linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select, face_resolution, apply_bg, apply_face]
406
+ diffusion_button.click(fn=stage2_process, inputs=stage2_ips, outputs=[result_gallery, event_id, fb_score, fb_text, face_gallery])
407
+ restart_button.click(fn=load_and_reset, inputs=[param_setting],
408
+ outputs=[edm_steps, s_cfg, s_stage2, s_stage1, s_churn, s_noise, a_prompt, n_prompt,
409
+ color_fix_type, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2])
410
+ submit_button.click(fn=submit_feedback, inputs=[event_id, fb_score, fb_text], outputs=[fb_text])
411
+ block.launch(server_name=server_ip, server_port=server_port)
gradio_demo_tiled.py ADDED
@@ -0,0 +1,339 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ import gradio as gr
4
+ from gradio_imageslider import ImageSlider
5
+ import argparse
6
+ from SUPIR.util import HWC3, upscale_image, fix_resize, convert_dtype
7
+ import numpy as np
8
+ import torch
9
+ from SUPIR.util import create_SUPIR_model, load_QF_ckpt
10
+ from PIL import Image
11
+ from llava.llava_agent import LLavaAgent
12
+ from CKPT_PTH import LLAVA_MODEL_PATH
13
+ import einops
14
+ import copy
15
+ import time
16
+ from omegaconf import OmegaConf
17
+ from sgm.modules.diffusionmodules.sampling import _sliding_windows
18
+
19
+ parser = argparse.ArgumentParser()
20
+ parser.add_argument("--ip", type=str, default='127.0.0.1')
21
+ parser.add_argument("--port", type=int, default='6688')
22
+ parser.add_argument("--no_llava", action='store_true', default=False)
23
+ parser.add_argument("--use_image_slider", action='store_true', default=False)
24
+ parser.add_argument("--log_history", action='store_true', default=False)
25
+ parser.add_argument("--loading_half_params", action='store_true', default=False)
26
+ parser.add_argument("--use_tile_vae", action='store_true', default=False)
27
+ parser.add_argument("--encoder_tile_size", type=int, default=512)
28
+ parser.add_argument("--decoder_tile_size", type=int, default=64)
29
+ parser.add_argument("--load_8bit_llava", action='store_true', default=False)
30
+ parser.add_argument("--local_prompt", action='store_true', default=False)
31
+ args = parser.parse_args()
32
+ server_ip = args.ip
33
+ server_port = args.port
34
+ use_llava = not args.no_llava
35
+
36
+ if torch.cuda.device_count() >= 2:
37
+ SUPIR_device = 'cuda:0'
38
+ LLaVA_device = 'cuda:1'
39
+ elif torch.cuda.device_count() == 1:
40
+ SUPIR_device = 'cuda:0'
41
+ LLaVA_device = 'cuda:0'
42
+ else:
43
+ raise ValueError('Currently support CUDA only.')
44
+
45
+ # load SUPIR
46
+ config_path = 'options/SUPIR_v0_tiled.yaml'
47
+ config = OmegaConf.load(config_path)
48
+ model = create_SUPIR_model(config_path, SUPIR_sign='Q')
49
+ if args.loading_half_params:
50
+ model = model.half()
51
+ if args.use_tile_vae:
52
+ model.init_tile_vae(encoder_tile_size=512, decoder_tile_size=64)
53
+ model = model.to(SUPIR_device)
54
+ model.first_stage_model.denoise_encoder_s1 = copy.deepcopy(model.first_stage_model.denoise_encoder)
55
+ model.current_model = 'v0-Q'
56
+ ckpt_Q, ckpt_F = load_QF_ckpt('options/SUPIR_v0.yaml')
57
+
58
+ tile_size = config.model.params.sampler_config.params.tile_size * 8
59
+ tile_stride = config.model.params.sampler_config.params.tile_stride * 8
60
+
61
+ # load LLaVA
62
+ if use_llava:
63
+ llava_agent = LLavaAgent(LLAVA_MODEL_PATH, device=LLaVA_device, load_8bit=args.load_8bit_llava, load_4bit=False)
64
+ else:
65
+ llava_agent = None
66
+
67
+ # only exhibit the overall quality of the stage1 output
68
+ def stage1_process(input_image, gamma_correction):
69
+ torch.cuda.set_device(SUPIR_device)
70
+ LQ = HWC3(input_image)
71
+ LQ = fix_resize(LQ, 512)
72
+ # stage1
73
+ LQ = np.array(LQ) / 255 * 2 - 1
74
+ LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
75
+ LQ = model.batchify_denoise(LQ, is_stage1=True)
76
+ LQ = (LQ[0].permute(1, 2, 0) * 127.5 + 127.5).cpu().numpy().round().clip(0, 255).astype(np.uint8)
77
+ # gamma correction
78
+ LQ = LQ / 255.0
79
+ LQ = np.power(LQ, gamma_correction)
80
+ LQ *= 255.0
81
+ LQ = LQ.round().clip(0, 255).astype(np.uint8)
82
+ return LQ
83
+
84
+ def llave_process(input_image, upscale, temperature, top_p, qs=None):
85
+ torch.cuda.set_device(SUPIR_device)
86
+ input_image = HWC3(input_image)
87
+ input_image = upscale_image(input_image, upscale, unit_resolution=32,
88
+ min_size=1024)
89
+ LQ = np.array(input_image) / 255.0
90
+ LQ *= 255.0
91
+ LQ = LQ.round().clip(0, 255).astype(np.uint8)
92
+ LQ = LQ / 255 * 2 - 1
93
+ LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
94
+ LQ = model.batchify_denoise(LQ, is_stage1=True)
95
+
96
+ _, _, h, w = LQ.shape
97
+ tiles_iterator = _sliding_windows(h, w, tile_size, tile_stride)
98
+ LQ_tiles = []
99
+ for hi, hi_end, wi, wi_end in tiles_iterator:
100
+ _LQ = LQ[:, :, hi:hi_end, wi:wi_end]
101
+ _LQ = (_LQ[0].permute(1, 2, 0) * 127.5 + 127.5).cpu().numpy().round().clip(0, 255).astype(np.uint8)
102
+ LQ_tiles.append(Image.fromarray(_LQ))
103
+
104
+ captions = []
105
+ torch.cuda.set_device(LLaVA_device)
106
+ if use_llava:
107
+ for LQ_tile in LQ_tiles:
108
+ captions += llava_agent.gen_image_caption([LQ_tile], temperature=temperature, top_p=top_p, qs=qs)
109
+ else:
110
+ captions = 'LLaVA is not available. Please add text manually.'
111
+ return str(captions)
112
+
113
+
114
+ def stage2_process(input_image, prompt, a_prompt, n_prompt, num_samples, upscale, edm_steps, s_stage1, s_stage2,
115
+ s_cfg, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction,
116
+ linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select):
117
+ torch.cuda.set_device(SUPIR_device)
118
+ event_id = str(time.time_ns())
119
+ event_dict = {'event_id': event_id, 'localtime': time.ctime(), 'prompt': prompt, 'a_prompt': a_prompt,
120
+ 'n_prompt': n_prompt, 'num_samples': num_samples, 'upscale': upscale, 'edm_steps': edm_steps,
121
+ 's_stage1': s_stage1, 's_stage2': s_stage2, 's_cfg': s_cfg, 'seed': seed, 's_churn': s_churn,
122
+ 's_noise': s_noise, 'color_fix_type': color_fix_type, 'diff_dtype': diff_dtype, 'ae_dtype': ae_dtype,
123
+ 'gamma_correction': gamma_correction, 'linear_CFG': linear_CFG, 'linear_s_stage2': linear_s_stage2,
124
+ 'spt_linear_CFG': spt_linear_CFG, 'spt_linear_s_stage2': spt_linear_s_stage2,
125
+ 'model_select': model_select}
126
+
127
+ if model_select != model.current_model:
128
+ if model_select == 'v0-Q':
129
+ print('load v0-Q')
130
+ model.load_state_dict(ckpt_Q, strict=False)
131
+ model.current_model = 'v0-Q'
132
+ elif model_select == 'v0-F':
133
+ print('load v0-F')
134
+ model.load_state_dict(ckpt_F, strict=False)
135
+ model.current_model = 'v0-F'
136
+ input_image = HWC3(input_image)
137
+ input_image = upscale_image(input_image, upscale, unit_resolution=32,
138
+ min_size=1024)
139
+
140
+ LQ = np.array(input_image) / 255.0
141
+ LQ = np.power(LQ, gamma_correction)
142
+ LQ *= 255.0
143
+ LQ = LQ.round().clip(0, 255).astype(np.uint8)
144
+ LQ = LQ / 255 * 2 - 1
145
+ LQ = torch.tensor(LQ, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
146
+ if use_llava:
147
+ captions = [eval(prompt)]
148
+ else:
149
+ captions = ['']
150
+
151
+ model.ae_dtype = convert_dtype(ae_dtype)
152
+ model.model.dtype = convert_dtype(diff_dtype)
153
+
154
+ samples = model.batchify_sample(LQ, captions, num_steps=edm_steps, restoration_scale=s_stage1, s_churn=s_churn,
155
+ s_noise=s_noise, cfg_scale=s_cfg, control_scale=s_stage2, seed=seed,
156
+ num_samples=num_samples, p_p=a_prompt, n_p=n_prompt, color_fix_type=color_fix_type,
157
+ use_linear_CFG=linear_CFG, use_linear_control_scale=linear_s_stage2,
158
+ cfg_scale_start=spt_linear_CFG, control_scale_start=spt_linear_s_stage2)
159
+
160
+ x_samples = (einops.rearrange(samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().round().clip(
161
+ 0, 255).astype(np.uint8)
162
+ results = [x_samples[i] for i in range(num_samples)]
163
+
164
+ if args.log_history:
165
+ os.makedirs(f'./history/{event_id[:5]}/{event_id[5:]}', exist_ok=True)
166
+ with open(f'./history/{event_id[:5]}/{event_id[5:]}/logs.txt', 'w') as f:
167
+ f.write(str(event_dict))
168
+ f.close()
169
+ Image.fromarray(input_image).save(f'./history/{event_id[:5]}/{event_id[5:]}/LQ.png')
170
+ for i, result in enumerate(results):
171
+ Image.fromarray(result).save(f'./history/{event_id[:5]}/{event_id[5:]}/HQ_{i}.png')
172
+ return [input_image] + results, event_id, 3, ''
173
+
174
+ def load_and_reset(param_setting):
175
+ edm_steps = 50
176
+ s_stage2 = 1.0
177
+ s_stage1 = -1.0
178
+ s_churn = 5
179
+ s_noise = 1.003
180
+ a_prompt = 'Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - ' \
181
+ 'realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore ' \
182
+ 'detailing, hyper sharpness, perfect without deformations.'
183
+ n_prompt = 'painting, oil painting, illustration, drawing, art, sketch, oil painting, cartoon, CG Style, ' \
184
+ '3D render, unreal engine, blurring, dirty, messy, worst quality, low quality, frames, watermark, ' \
185
+ 'signature, jpeg artifacts, deformed, lowres, over-smooth'
186
+ color_fix_type = 'Wavelet'
187
+ spt_linear_s_stage2 = 0.0
188
+ linear_s_stage2 = False
189
+ linear_CFG = True
190
+ if param_setting == "Quality":
191
+ s_cfg = 7.5
192
+ spt_linear_CFG = 4.0
193
+ elif param_setting == "Fidelity":
194
+ s_cfg = 4.0
195
+ spt_linear_CFG = 1.0
196
+ else:
197
+ raise NotImplementedError
198
+ return edm_steps, s_cfg, s_stage2, s_stage1, s_churn, s_noise, a_prompt, n_prompt, color_fix_type, linear_CFG, \
199
+ linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2
200
+
201
+
202
+ def submit_feedback(event_id, fb_score, fb_text):
203
+ if args.log_history:
204
+ with open(f'./history/{event_id[:5]}/{event_id[5:]}/logs.txt', 'r') as f:
205
+ event_dict = eval(f.read())
206
+ f.close()
207
+ event_dict['feedback'] = {'score': fb_score, 'text': fb_text}
208
+ with open(f'./history/{event_id[:5]}/{event_id[5:]}/logs.txt', 'w') as f:
209
+ f.write(str(event_dict))
210
+ f.close()
211
+ return 'Submit successfully, thank you for your comments!'
212
+ else:
213
+ return 'Submit failed, the server is not set to log history.'
214
+
215
+ title_md = """
216
+ # **SUPIR: Practicing Model Scaling for Photo-Realistic Image Restoration**
217
+
218
+ ⚠️SUPIR is still a research project under tested and is not yet a stable commercial product.
219
+
220
+ [[Paper](https://arxiv.org/abs/2401.13627)] &emsp; [[Project Page](http://supir.xpixel.group/)] &emsp; [[How to play](https://github.com/Fanghua-Yu/SUPIR/blob/master/assets/DemoGuide.png)]
221
+ """
222
+
223
+
224
+ claim_md = """
225
+ ## **Terms of use**
226
+
227
+ By using this service, users are required to agree to the following terms: The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research. Please submit a feedback to us if you get any inappropriate answer! We will collect those to keep improving our models. For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality.
228
+
229
+ ## **License**
230
+
231
+ The service is a research preview intended for non-commercial use only, subject to the model [License](https://github.com/Fanghua-Yu/SUPIR) of SUPIR.
232
+ """
233
+
234
+
235
+ block = gr.Blocks(title='SUPIR').queue()
236
+ with block:
237
+ with gr.Row():
238
+ gr.Markdown(title_md)
239
+ with gr.Row():
240
+ with gr.Column():
241
+ with gr.Row(equal_height=True):
242
+ with gr.Column():
243
+ gr.Markdown("<center>Input</center>")
244
+ input_image = gr.Image(type="numpy", elem_id="image-input", height=400, width=400)
245
+ with gr.Column():
246
+ gr.Markdown("<center>Stage1 Output</center>")
247
+ denoise_image = gr.Image(type="numpy", elem_id="image-s1", height=400, width=400)
248
+ prompt = gr.Textbox(label="Prompt", value="")
249
+ with gr.Accordion("Stage1 options", open=False):
250
+ gamma_correction = gr.Slider(label="Gamma Correction", minimum=0.1, maximum=2.0, value=1.0, step=0.1)
251
+ with gr.Accordion("LLaVA options", open=False):
252
+ temperature = gr.Slider(label="Temperature", minimum=0., maximum=1.0, value=0.2, step=0.1)
253
+ top_p = gr.Slider(label="Top P", minimum=0., maximum=1.0, value=0.7, step=0.1)
254
+ qs = gr.Textbox(label="Question", value="Describe this image and its style in a very detailed manner. "
255
+ "The image is a realistic photography, not an art painting.")
256
+ with gr.Accordion("Stage2 options", open=False):
257
+ num_samples = gr.Slider(label="Num Samples", minimum=1, maximum=4 if not args.use_image_slider else 1
258
+ , value=1, step=1)
259
+ upscale = gr.Slider(label="Upscale", minimum=1, maximum=8, value=1, step=1)
260
+ edm_steps = gr.Slider(label="Steps", minimum=20, maximum=200, value=50, step=1)
261
+ s_cfg = gr.Slider(label="Text Guidance Scale", minimum=1.0, maximum=15.0, value=7.5, step=0.1)
262
+ s_stage2 = gr.Slider(label="Stage2 Guidance Strength", minimum=0., maximum=1., value=1., step=0.05)
263
+ s_stage1 = gr.Slider(label="Stage1 Guidance Strength", minimum=-1.0, maximum=6.0, value=-1.0, step=1.0)
264
+ seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=True)
265
+ s_churn = gr.Slider(label="S-Churn", minimum=0, maximum=40, value=5, step=1)
266
+ s_noise = gr.Slider(label="S-Noise", minimum=1.0, maximum=1.1, value=1.003, step=0.001)
267
+ a_prompt = gr.Textbox(label="Default Positive Prompt",
268
+ value='Cinematic, High Contrast, highly detailed, taken using a Canon EOS R '
269
+ 'camera, hyper detailed photo - realistic maximum detail, 32k, Color '
270
+ 'Grading, ultra HD, extreme meticulous detailing, skin pore detailing, '
271
+ 'hyper sharpness, perfect without deformations.')
272
+ n_prompt = gr.Textbox(label="Default Negative Prompt",
273
+ value='painting, oil painting, illustration, drawing, art, sketch, oil painting, '
274
+ 'cartoon, CG Style, 3D render, unreal engine, blurring, dirty, messy, '
275
+ 'worst quality, low quality, frames, watermark, signature, jpeg artifacts, '
276
+ 'deformed, lowres, over-smooth')
277
+ with gr.Row():
278
+ with gr.Column():
279
+ linear_CFG = gr.Checkbox(label="Linear CFG", value=True)
280
+ spt_linear_CFG = gr.Slider(label="CFG Start", minimum=1.0,
281
+ maximum=9.0, value=4.0, step=0.5)
282
+ with gr.Column():
283
+ linear_s_stage2 = gr.Checkbox(label="Linear Stage2 Guidance", value=False)
284
+ spt_linear_s_stage2 = gr.Slider(label="Guidance Start", minimum=0.,
285
+ maximum=1., value=0., step=0.05)
286
+ with gr.Row():
287
+ with gr.Column():
288
+ diff_dtype = gr.Radio(['fp32', 'fp16', 'bf16'], label="Diffusion Data Type", value="fp16",
289
+ interactive=True)
290
+ with gr.Column():
291
+ ae_dtype = gr.Radio(['fp32', 'bf16'], label="Auto-Encoder Data Type", value="bf16",
292
+ interactive=True)
293
+ with gr.Column():
294
+ color_fix_type = gr.Radio(["None", "AdaIn", "Wavelet"], label="Color-Fix Type", value="Wavelet",
295
+ interactive=True)
296
+ with gr.Column():
297
+ model_select = gr.Radio(["v0-Q", "v0-F"], label="Model Selection", value="v0-Q",
298
+ interactive=True)
299
+
300
+ with gr.Column():
301
+ gr.Markdown("<center>Stage2 Output</center>")
302
+ if not args.use_image_slider:
303
+ result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery1")
304
+ else:
305
+ result_gallery = ImageSlider(label='Output', show_label=False, elem_id="gallery1")
306
+ with gr.Row():
307
+ with gr.Column():
308
+ denoise_button = gr.Button(value="Stage1 Run")
309
+ with gr.Column():
310
+ llave_button = gr.Button(value="LlaVa Run")
311
+ with gr.Column():
312
+ diffusion_button = gr.Button(value="Stage2 Run")
313
+ with gr.Row():
314
+ with gr.Column():
315
+ param_setting = gr.Dropdown(["Quality", "Fidelity"], interactive=True, label="Param Setting",
316
+ value="Quality")
317
+ with gr.Column():
318
+ restart_button = gr.Button(value="Reset Param", scale=2)
319
+ with gr.Accordion("Feedback", open=True):
320
+ fb_score = gr.Slider(label="Feedback Score", minimum=1, maximum=5, value=3, step=1,
321
+ interactive=True)
322
+ fb_text = gr.Textbox(label="Feedback Text", value="", placeholder='Please enter your feedback here.')
323
+ submit_button = gr.Button(value="Submit Feedback")
324
+ with gr.Row():
325
+ gr.Markdown(claim_md)
326
+ event_id = gr.Textbox(label="Event ID", value="", visible=False)
327
+
328
+ llave_button.click(fn=llave_process, inputs=[input_image, upscale, temperature, top_p, qs], outputs=[prompt])
329
+ denoise_button.click(fn=stage1_process, inputs=[input_image, gamma_correction],
330
+ outputs=[denoise_image])
331
+ stage2_ips = [input_image, prompt, a_prompt, n_prompt, num_samples, upscale, edm_steps, s_stage1, s_stage2,
332
+ s_cfg, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction,
333
+ linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select]
334
+ diffusion_button.click(fn=stage2_process, inputs=stage2_ips, outputs=[result_gallery, event_id, fb_score, fb_text])
335
+ restart_button.click(fn=load_and_reset, inputs=[param_setting],
336
+ outputs=[edm_steps, s_cfg, s_stage2, s_stage1, s_churn, s_noise, a_prompt, n_prompt,
337
+ color_fix_type, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2])
338
+ submit_button.click(fn=submit_feedback, inputs=[event_id, fb_score, fb_text], outputs=[fb_text])
339
+ block.launch(server_name=server_ip, server_port=server_port)
requirements.txt ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ fastapi==0.95.1
2
+ gradio==4.16.0
3
+ gradio_imageslider==0.0.17
4
+ gradio_client==0.8.1
5
+ Markdown==3.4.1
6
+ numpy==1.24.2
7
+ requests==2.28.2
8
+ sentencepiece==0.1.98
9
+ tokenizers==0.13.3
10
+ torch>=2.1.0
11
+ torchvision>=0.16.0
12
+ uvicorn==0.21.1
13
+ wandb==0.14.0
14
+ httpx==0.24.0
15
+ transformers==4.28.1
16
+ accelerate==0.18.0
17
+ scikit-learn==1.2.2
18
+ sentencepiece==0.1.98
19
+ einops==0.7.0
20
+ einops-exts==0.0.4
21
+ timm==0.9.8
22
+ openai-clip==1.0.1
23
+ fsspec==2023.4.0
24
+ kornia==0.6.9
25
+ matplotlib==3.7.1
26
+ ninja==1.11.1
27
+ omegaconf==2.3.0
28
+ open-clip-torch==2.17.1
29
+ opencv-python==4.7.0.72
30
+ pandas==2.0.1
31
+ Pillow==9.4.0
32
+ pytorch-lightning==2.1.2
33
+ PyYAML==6.0
34
+ scipy==1.9.1
35
+ tqdm==4.65.0
36
+ triton==2.1.0
37
+ urllib3==1.26.15
38
+ webdataset==0.2.48
39
+ xformers>=0.0.20
40
+ facexlib==0.3.0
41
+ k-diffusion==0.1.1.post1
42
+ diffusers==0.16.1
test.py ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch.cuda
2
+ import argparse
3
+ from SUPIR.util import create_SUPIR_model, PIL2Tensor, Tensor2PIL, convert_dtype
4
+ from PIL import Image
5
+ from llava.llava_agent import LLavaAgent
6
+ from CKPT_PTH import LLAVA_MODEL_PATH
7
+ import os
8
+ from torch.nn.functional import interpolate
9
+
10
+ if torch.cuda.device_count() >= 2:
11
+ SUPIR_device = 'cuda:0'
12
+ LLaVA_device = 'cuda:1'
13
+ elif torch.cuda.device_count() == 1:
14
+ SUPIR_device = 'cuda:0'
15
+ LLaVA_device = 'cuda:0'
16
+ else:
17
+ raise ValueError('Currently support CUDA only.')
18
+
19
+ # hyparams here
20
+ parser = argparse.ArgumentParser()
21
+ parser.add_argument("--img_dir", type=str)
22
+ parser.add_argument("--save_dir", type=str)
23
+ parser.add_argument("--upscale", type=int, default=1)
24
+ parser.add_argument("--SUPIR_sign", type=str, default='Q', choices=['F', 'Q'])
25
+ parser.add_argument("--seed", type=int, default=1234)
26
+ parser.add_argument("--min_size", type=int, default=1024)
27
+ parser.add_argument("--edm_steps", type=int, default=50)
28
+ parser.add_argument("--s_stage1", type=int, default=-1)
29
+ parser.add_argument("--s_churn", type=int, default=5)
30
+ parser.add_argument("--s_noise", type=float, default=1.003)
31
+ parser.add_argument("--s_cfg", type=float, default=7.5)
32
+ parser.add_argument("--s_stage2", type=float, default=1.)
33
+ parser.add_argument("--num_samples", type=int, default=1)
34
+ parser.add_argument("--a_prompt", type=str,
35
+ default='Cinematic, High Contrast, highly detailed, taken using a Canon EOS R '
36
+ 'camera, hyper detailed photo - realistic maximum detail, 32k, Color '
37
+ 'Grading, ultra HD, extreme meticulous detailing, skin pore detailing, '
38
+ 'hyper sharpness, perfect without deformations.')
39
+ parser.add_argument("--n_prompt", type=str,
40
+ default='painting, oil painting, illustration, drawing, art, sketch, oil painting, '
41
+ 'cartoon, CG Style, 3D render, unreal engine, blurring, dirty, messy, '
42
+ 'worst quality, low quality, frames, watermark, signature, jpeg artifacts, '
43
+ 'deformed, lowres, over-smooth')
44
+ parser.add_argument("--color_fix_type", type=str, default='Wavelet', choices=["None", "AdaIn", "Wavelet"])
45
+ parser.add_argument("--linear_CFG", action='store_true', default=True)
46
+ parser.add_argument("--linear_s_stage2", action='store_true', default=False)
47
+ parser.add_argument("--spt_linear_CFG", type=float, default=4.0)
48
+ parser.add_argument("--spt_linear_s_stage2", type=float, default=0.)
49
+ parser.add_argument("--ae_dtype", type=str, default="bf16", choices=['fp32', 'bf16'])
50
+ parser.add_argument("--diff_dtype", type=str, default="fp16", choices=['fp32', 'fp16', 'bf16'])
51
+ parser.add_argument("--no_llava", action='store_true', default=False)
52
+ parser.add_argument("--loading_half_params", action='store_true', default=False)
53
+ parser.add_argument("--use_tile_vae", action='store_true', default=False)
54
+ parser.add_argument("--encoder_tile_size", type=int, default=512)
55
+ parser.add_argument("--decoder_tile_size", type=int, default=64)
56
+ parser.add_argument("--load_8bit_llava", action='store_true', default=False)
57
+ args = parser.parse_args()
58
+ print(args)
59
+ use_llava = not args.no_llava
60
+
61
+ # load SUPIR
62
+ model = create_SUPIR_model('options/SUPIR_v0.yaml', SUPIR_sign=args.SUPIR_sign)
63
+ if args.loading_half_params:
64
+ model = model.half()
65
+ if args.use_tile_vae:
66
+ model.init_tile_vae(encoder_tile_size=args.encoder_tile_size, decoder_tile_size=args.decoder_tile_size)
67
+ model.ae_dtype = convert_dtype(args.ae_dtype)
68
+ model.model.dtype = convert_dtype(args.diff_dtype)
69
+ model = model.to(SUPIR_device)
70
+ # load LLaVA
71
+ if use_llava:
72
+ llava_agent = LLavaAgent(LLAVA_MODEL_PATH, device=LLaVA_device, load_8bit=args.load_8bit_llava, load_4bit=False)
73
+ else:
74
+ llava_agent = None
75
+
76
+ os.makedirs(args.save_dir, exist_ok=True)
77
+ for img_pth in os.listdir(args.img_dir):
78
+ img_name = os.path.splitext(img_pth)[0]
79
+
80
+ LQ_ips = Image.open(os.path.join(args.img_dir, img_pth))
81
+ LQ_img, h0, w0 = PIL2Tensor(LQ_ips, upsacle=args.upscale, min_size=args.min_size)
82
+ LQ_img = LQ_img.unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
83
+
84
+ # step 1: Pre-denoise for LLaVA, resize to 512
85
+ LQ_img_512, h1, w1 = PIL2Tensor(LQ_ips, upsacle=args.upscale, min_size=args.min_size, fix_resize=512)
86
+ LQ_img_512 = LQ_img_512.unsqueeze(0).to(SUPIR_device)[:, :3, :, :]
87
+ clean_imgs = model.batchify_denoise(LQ_img_512)
88
+ clean_PIL_img = Tensor2PIL(clean_imgs[0], h1, w1)
89
+
90
+ # step 2: LLaVA
91
+ if use_llava:
92
+ captions = llava_agent.gen_image_caption([clean_PIL_img])
93
+ else:
94
+ captions = ['']
95
+ print(captions)
96
+
97
+ # # step 3: Diffusion Process
98
+ samples = model.batchify_sample(LQ_img, captions, num_steps=args.edm_steps, restoration_scale=args.s_stage1, s_churn=args.s_churn,
99
+ s_noise=args.s_noise, cfg_scale=args.s_cfg, control_scale=args.s_stage2, seed=args.seed,
100
+ num_samples=args.num_samples, p_p=args.a_prompt, n_p=args.n_prompt, color_fix_type=args.color_fix_type,
101
+ use_linear_CFG=args.linear_CFG, use_linear_control_scale=args.linear_s_stage2,
102
+ cfg_scale_start=args.spt_linear_CFG, control_scale_start=args.spt_linear_s_stage2)
103
+ # save
104
+ for _i, sample in enumerate(samples):
105
+ Tensor2PIL(sample, h0, w0).save(f'{args.save_dir}/{img_name}_{_i}.png')
106
+