Zongsheng commited on
Commit
b2aaa70
1 Parent(s): cd4fbd5

first upload

Browse files
app.py ADDED
@@ -0,0 +1,139 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+ # -*- coding:utf-8 -*-
3
+ # Power by Zongsheng Yue 2022-12-16 16:17:14
4
+
5
+ import os
6
+ import torch
7
+ import argparse
8
+ import numpy as np
9
+ import grdio as gr
10
+ from pathlib import Path
11
+ from einops import rearrange
12
+ from omegaconf import OmegaConf
13
+ from skimage import img_as_ubyte
14
+
15
+ from utils import util_opts
16
+ from utils import util_image
17
+ from utils import util_common
18
+
19
+ from sampler import DifIRSampler
20
+ from ResizeRight.resize_right import resize
21
+ from basicsr.utils.download_util import load_file_from_url
22
+
23
+ def predict(im_path, background_enhance, face_upsample, upscale, started_timesteps):
24
+ cfg_path = 'configs/sample/iddpm_ffhq512_swinir.yaml'
25
+
26
+ # setting configurations
27
+ configs = OmegaConf.load(cfg_path)
28
+ configs.aligned = False
29
+ configs.background_enhance = background_enhance
30
+ configs.face_upsample = face_upsample
31
+
32
+ started_timesteps = int(started_timesteps)
33
+ assert started_timesteps < int(configs.diffusion.params.timestep_respacing)
34
+
35
+ # prepare the checkpoint
36
+ if not Path(configs.model.ckpt_path).exists():
37
+ load_file_from_url(
38
+ url="https://github.com/zsyOAOA/DifFace/releases/download/V1.0/iddpm_ffhq512_ema500000.pth",
39
+ model_dir=str(Path(configs.model.ckpt_path).parent),
40
+ progress=True,
41
+ file_name=Path(configs.model.ckpt_path).name,
42
+ )
43
+ if not Path(configs.model_ir.ckpt_path).exists():
44
+ load_file_from_url(
45
+ url="https://github.com/zsyOAOA/DifFace/releases/download/V1.0/General_Face_ffhq512.pth",
46
+ model_dir=str(Path(configs.model_ir.ckpt_path).parent),
47
+ progress=True,
48
+ file_name=Path(configs.model_ir.ckpt_path).name,
49
+ )
50
+
51
+ # Load image
52
+ im_lq = util_image.imread(im_path, chn='bgr', dtype='uint8')
53
+ if upscale > 4:
54
+ upscale = 4 # avoid momory exceeded due to too large upscale
55
+ if upscale > 2 and min(im_lq.shape[:2])>1280:
56
+ upscale = 2 # avoid momory exceeded due to too large img resolution
57
+ configs.detection.upscale = int(upscale)
58
+
59
+ # build the sampler for diffusion
60
+ sampler_dist = DifIRSampler(configs)
61
+
62
+ image_restored, face_restored, face_cropped = sampler_dist.sample_func_bfr_unaligned(
63
+ y0=im_lq,
64
+ start_timesteps=started_timesteps,
65
+ need_restoration=True,
66
+ draw_box=False,
67
+ )
68
+
69
+ restored_image_dir = Path('restored_output')
70
+ if not restored_image_dir.exists():
71
+ restored_image_dir.mkdir()
72
+ # save the whole image
73
+ save_path = restored_image_dir / Path(im_path).name
74
+ util_image.imwrite(image_restored, save_path, chn='bgr', dtype_in='uint8')
75
+
76
+ return image_restored, str(save_path)
77
+
78
+ # im_path = './testdata/whole_imgs/00.jpg'
79
+ # predict(im_path, True, True, 3, 100)
80
+
81
+ title = "DifFace: Blind Face Restoration with Diffused Error Contraction"
82
+ description = r"""<center><img src='./assets/DifFace_Framework.png' alt='DifFace logo'></center>
83
+ <b>Official Gradio demo</b> for <a href='https://github.com/zsyOAOA/DifFace' target='_blank'><b>DifFace: Blind Face Restoration with Diffused Error Contraction</b></a>.<br>
84
+ 🔥 DifFace is a robust face restoration algorithm for old or corrupted photos.<br>
85
+ """
86
+ article = r"""
87
+ If DifFace is helpful for your work, please help to ⭐ the <a href='https://github.com/zsyOAOA/DifFace' target='_blank'>Github Repo</a>. Thanks!
88
+ [![GitHub Stars](https://img.shields.io/github/stars/zsyOAOA/DifFace?affiliations=OWNER&color=green&style=social)](https://github.com/zsyOAOA/DifFace)
89
+
90
+ ---
91
+
92
+ 📝 **Citation**
93
+
94
+ If our work is useful for your research, please consider citing:
95
+ ```bibtex
96
+ @article{yue2022difface,
97
+ title={DifFace: Blind Face Restoration with Diffused Error Contraction},
98
+ author={Yue, Zongsheng and Loy, Chen Change},
99
+ journal={arXiv preprint arXiv:2212.06512},
100
+ year={2022}
101
+ }
102
+ ```
103
+
104
+ 📋 **License**
105
+
106
+ This project is licensed under <a rel="license" href="https://github.com/zsyOAOA/DifFace/blob/master/LICENSE">S-Lab License 1.0</a>.
107
+ Redistribution and use for non-commercial purposes should follow this license.
108
+
109
+ 📧 **Contact**
110
+ If you have any questions, please feel free to contact me via <b>zsyzam@gmail.com</b>.
111
+ ![visitors](https://visitor-badge.laobi.icu/badge?page_id=zsyOAOA/DifFace)
112
+ """
113
+
114
+ demo = gr.Interface(
115
+ inference,
116
+ inputs=[
117
+ gr.inputs.Image(type="filepath", label="Input"),
118
+ gr.inputs.Checkbox(default=True, label="Background_Enhance"),
119
+ gr.inputs.Checkbox(default=True, label="Face_Upsample"),
120
+ gr.inputs.Number(default=2, label="Rescaling_Factor (up to 4)"),
121
+ gr.Slider(1, 200, value=100, step=10, label='Realism-Fidelity Trade-off')
122
+ ],
123
+ outputs=[
124
+ gr.outputs.Image(type="numpy", label="Output"),
125
+ gr.outputs.File(label="Download the output")
126
+ ],
127
+ title=title,
128
+ description=description,
129
+ article=article,
130
+ examples=[
131
+ ['./testdata/whole_imgs/00.jpg', True, True, 2, 100],
132
+ ['./testdata/whole_imgs/01.jpg', True, True, 2, 100],
133
+ ['./testdata/whole_imgs/04.jpg', True, True, 2, 100],
134
+ ['./testdata/whole_imgs/Solvay_conference_1927.png', True, True, 2, 100],
135
+ ]
136
+ )
137
+
138
+ demo.queue(concurrency_count=4)
139
+ demo.launch()
configs/.DS_Store ADDED
Binary file (6.15 kB). View file
 
configs/sample/iddpm_ffhq512.yaml ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ gpu_id: ""
2
+ seed: 10000
3
+ display: True
4
+ im_size: 512
5
+
6
+ diffusion:
7
+ target: models.script_util.create_gaussian_diffusion
8
+ params:
9
+ steps: 1000
10
+ learn_sigma: True
11
+ sigma_small: False
12
+ noise_schedule: linear
13
+ use_kl: False
14
+ predict_xstart: False
15
+ rescale_timesteps: False
16
+ rescale_learned_sigmas: True
17
+ timestep_respacing: "1000"
18
+
19
+ model:
20
+ target: models.unet.UNetModel
21
+ ckpt_path: pretrained_zoo/iddpm_ffhq512/ema0999_model_500000.pth
22
+ params:
23
+ image_size: 512
24
+ in_channels: 3
25
+ model_channels: 32
26
+ out_channels: 6
27
+ attention_resolutions: [32, 16, 8]
28
+ dropout: 0
29
+ channel_mult: [1, 2, 4, 8, 8, 16, 16]
30
+ num_res_blocks: [1, 2, 2, 2, 2, 3, 4]
31
+ conv_resample: True
32
+ dims: 2
33
+ use_fp16: False
34
+ num_head_channels: 64
35
+ use_scale_shift_norm: True
36
+ resblock_updown: False
37
+ use_new_attention_order: False
38
+
39
+ model_ir: ~
configs/sample/iddpm_ffhq512_swinir.yaml ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ gpu_id: ""
2
+ seed: 10000
3
+ display: True
4
+ im_size: 512
5
+ aligned: True
6
+ background_enhance: True
7
+ face_upsample: True
8
+
9
+ diffusion:
10
+ target: models.script_util.create_gaussian_diffusion
11
+ params:
12
+ steps: 1000
13
+ learn_sigma: True
14
+ sigma_small: False
15
+ noise_schedule: linear
16
+ use_kl: False
17
+ predict_xstart: False
18
+ rescale_timesteps: False
19
+ rescale_learned_sigmas: True
20
+ timestep_respacing: "250"
21
+
22
+ model:
23
+ target: models.unet.UNetModel
24
+ ckpt_path: ./weights/diffusion/iddpm_ffhq512_ema500000.pth
25
+ params:
26
+ image_size: 512
27
+ in_channels: 3
28
+ model_channels: 32
29
+ out_channels: 6
30
+ attention_resolutions: [32, 16, 8]
31
+ dropout: 0
32
+ channel_mult: [1, 2, 4, 8, 8, 16, 16]
33
+ num_res_blocks: [1, 2, 2, 2, 2, 3, 4]
34
+ conv_resample: True
35
+ dims: 2
36
+ use_fp16: False
37
+ num_head_channels: 64
38
+ use_scale_shift_norm: True
39
+ resblock_updown: False
40
+ use_new_attention_order: False
41
+
42
+ model_ir:
43
+ target: models.swinir.SwinIR
44
+ ckpt_path: ./weights/SwinIR/General_Face_ffhq512.pth
45
+ params:
46
+ img_size: 64
47
+ patch_size: 1
48
+ in_chans: 3
49
+ embed_dim: 180
50
+ depths: [6, 6, 6, 6, 6, 6, 6, 6]
51
+ num_heads: [6, 6, 6, 6, 6, 6, 6, 6]
52
+ window_size: 8
53
+ mlp_ratio: 2
54
+ sf: 8
55
+ img_range: 1.0
56
+ upsampler: "nearest+conv"
57
+ resi_connection: "1conv"
58
+ unshuffle: True
59
+ unshuffle_scale: 8
60
+
61
+ # face detection model for unaligned face
62
+ detection:
63
+ det_model: "YOLOv5l" # large model: 'YOLOv5l', 'retinaface_resnet50'; small model: 'YOLOv5n', 'retinaface_mobile0.25'
64
+ upscale: 2 # The final upscaling factor for the whole image
configs/training/diffusion_ffhq512.yaml ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ model:
2
+ target: models.unet.UNetModel
3
+ params:
4
+ image_size: 512
5
+ in_channels: 3
6
+ model_channels: 32
7
+ out_channels: 6
8
+ attention_resolutions: [32, 16, 8]
9
+ dropout: 0
10
+ channel_mult: [1, 2, 4, 8, 8, 16, 16]
11
+ num_res_blocks: [1, 2, 2, 2, 2, 3, 4]
12
+ conv_resample: True
13
+ dims: 2
14
+ use_fp16: False
15
+ num_head_channels: 64
16
+ use_scale_shift_norm: True
17
+ resblock_updown: False
18
+ use_new_attention_order: False
19
+
20
+ diffusion:
21
+ target: models.script_util.create_gaussian_diffusion
22
+ params:
23
+ steps: 1000
24
+ learn_sigma: True
25
+ sigma_small: False
26
+ noise_schedule: linear
27
+ use_kl: False
28
+ predict_xstart: False
29
+ rescale_timesteps: False
30
+ rescale_learned_sigmas: True
31
+ timestep_respacing: ""
32
+
33
+ train:
34
+ lr: 1e-4
35
+ batch: [32, 4] # batchsize for training and validation
36
+ microbatch: 8
37
+ use_fp16: False
38
+ num_workers: 16
39
+ prefetch_factor: 2
40
+ iterations: 800000
41
+ weight_decay: 0
42
+ scheduler: step # step or cosin
43
+ milestones: [10000, 800000]
44
+ ema_rates: [0.999]
45
+ save_freq: 10000
46
+ val_freq: 5000
47
+ log_freq: [1000, 2000]
48
+
49
+ data:
50
+ train:
51
+ type: face
52
+ params:
53
+ ffhq_txt: ./datapipe/files_txt/ffhq512.txt
54
+ out_size: 512
55
+ transform_type: face
configs/training/swinir_ffhq512.yaml ADDED
@@ -0,0 +1,67 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ model:
2
+ target: models.swinir.SwinIR
3
+ params:
4
+ img_size: 64
5
+ patch_size: 1
6
+ in_chans: 3
7
+ embed_dim: 180
8
+ depths: [6, 6, 6, 6, 6, 6, 6, 6]
9
+ num_heads: [6, 6, 6, 6, 6, 6, 6, 6]
10
+ window_size: 8
11
+ mlp_ratio: 2
12
+ sf: 8
13
+ img_range: 1.0
14
+ upsampler: "nearest+conv"
15
+ resi_connection: "1conv"
16
+ unshuffle: True
17
+ unshuffle_scale: 8
18
+
19
+ train:
20
+ lr: 1e-4
21
+ lr_min: 5e-6
22
+ batch: [16, 4] # batchsize for training and validation
23
+ microbatch: 2
24
+ num_workers: 8
25
+ prefetch_factor: 2
26
+ iterations: 800000
27
+ weight_decay: 0
28
+ save_freq: 20000
29
+ val_freq: 20000
30
+ log_freq: [100, 2000, 50]
31
+
32
+ data:
33
+ train:
34
+ type: gfpgan
35
+ params:
36
+ files_txt: ./datapipe/files_txt/ffhq512_train.txt
37
+ io_backend:
38
+ type: disk
39
+
40
+ use_hflip: true
41
+ mean: [0.0, 0.0, 0.0]
42
+ std: [1.0, 1.0, 1.0]
43
+ out_size: 512
44
+
45
+ blur_kernel_size: 41
46
+ kernel_list: ['iso', 'aniso']
47
+ kernel_prob: [0.5, 0.5]
48
+ blur_sigma: [0.1, 15]
49
+ downsample_range: [0.8, 32]
50
+ noise_range: [0, 20]
51
+ jpeg_range: [30, 100]
52
+
53
+ color_jitter_prob: ~
54
+ color_jitter_pt_prob: ~
55
+ gray_prob: 0.01
56
+ gt_gray: True
57
+
58
+ need_gt_path: False
59
+ val:
60
+ type: folder
61
+ params:
62
+ dir_path: /mnt/vdb/IRDiff/Face/testing_data/celeba512_lq
63
+ dir_path_gt: /mnt/vdb/IRDiff/Face/testing_data/celeba512_hq
64
+ ext: png
65
+ need_gt_path: False
66
+ length: ~
67
+
testdata/.DS_Store ADDED
Binary file (6.15 kB). View file
 
testdata/whole_imgs/.DS_Store ADDED
Binary file (6.15 kB). View file
 
testdata/whole_imgs/00.jpg ADDED
testdata/whole_imgs/01.jpg ADDED
testdata/whole_imgs/02.png ADDED
testdata/whole_imgs/03.png ADDED
testdata/whole_imgs/04.jpg ADDED
testdata/whole_imgs/05.jpg ADDED
testdata/whole_imgs/Solvay_conference_1927.png ADDED