avans06 commited on
Commit
debfe4b
1 Parent(s): c85d0ce

Overall improvements with added support for various face restoration models and upscale models.

Browse files

01. The inference uses the auto_split_upscale mechanism, auto_split_upscale is in the dataops.py file, and its source is from the ESRGAN project forked by authors joeyballentine and BlueAmulet.

02. The face model now supports RestoreFormer++, authored by wzhouxiff.

03. The face model now supports CodeFormer, authored by sczhou.

04. The face model now supports GPEN, authored by yangxy.

05. Added support for parsing older RRDB models from the ESRGAN project.

06. Added support for parsing DAT models from the DAT project by author zhengchen1999.

07. Added support for parsing HAT models from the HAT project by author XPixelGroup.

08. Added support for parsing RealPLKSR models from the PLKSR project by author dslisleedh & neosr-project.

09. In the interface, you can select the type of face detection: "retinaface_resnet50", "YOLOv5l", or "YOLOv5n". The default is "retinaface_resnet50".

10. In the interface, you can set whether to include the model name in the output image file names for easier result verification.

11. The Gradio interface has been updated from the `Interface` syntax to the `Blocks` syntax.

12. Upgrade the project's Gradio version to 5.9.0.

.gitignore CHANGED
@@ -140,4 +140,5 @@ dmypy.json
140
  .vs
141
  output
142
  weights
143
- *.jpg
 
 
140
  .vs
141
  output
142
  weights
143
+ .jpg
144
+ .png
README.md CHANGED
@@ -1,12 +1,12 @@
1
  ---
2
- title: Image Face Upscale Restoration-GFPGAN
3
  emoji: 📈
4
  colorFrom: blue
5
  colorTo: gray
6
  sdk: gradio
7
- sdk_version: 5.8.0
8
  app_file: app.py
9
- pinned: false
10
  license: apache-2.0
11
  ---
12
 
 
1
  ---
2
+ title: Image Face Upscale Restoration-GFPGAN-RestoreFormerPlusPlus-CodeFormer
3
  emoji: 📈
4
  colorFrom: blue
5
  colorTo: gray
6
  sdk: gradio
7
+ sdk_version: 5.9.0
8
  app_file: app.py
9
+ pinned: true
10
  license: apache-2.0
11
  ---
12
 
app.py CHANGED
@@ -1,7 +1,6 @@
1
  import os
2
  import gc
3
  import cv2
4
- import requests
5
  import numpy as np
6
  import gradio as gr
7
  import torch
@@ -11,59 +10,208 @@ from realesrgan.utils import RealESRGANer
11
 
12
 
13
  # Define URLs and their corresponding local storage paths
14
- face_model = {
15
- "GFPGANv1.4.pth": "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth",
16
- "RestoreFormer++.ckpt": "https://github.com/wzhouxiff/RestoreFormerPlusPlus/releases/download/v1.0.0/RestoreFormer++.ckpt",
17
- # "CodeFormer.pth": "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/CodeFormer.pth",
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
18
  # legacy model
19
- "GFPGANv1.3.pth": "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth",
20
- "GFPGANv1.2.pth": "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.2.pth",
21
- "RestoreFormer.ckpt": "https://github.com/wzhouxiff/RestoreFormerPlusPlus/releases/download/v1.0.0/RestoreFormer.ckpt",
 
 
 
22
  }
23
- realesr_model = {
24
  # SRVGGNet
25
- "realesr-general-x4v3.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth", # x4 SRVGGNet (S size)
26
- "realesr-animevideov3.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth", # x4 SRVGGNet (XS size)
 
 
 
 
 
 
27
  # RRDBNet
28
- "RealESRGAN_x4plus_anime_6B.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth", # x4 RRDBNet with 6 blocks
29
- "RealESRGAN_x2plus.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
30
- "RealESRNet_x4plus.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth",
31
- "RealESRGAN_x4plus.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
 
 
 
 
 
 
 
 
 
 
 
 
32
  # ESRGAN(oldRRDB)
33
- "4x-AnimeSharp.pth": "https://huggingface.co/utnah/esrgan/resolve/main/4x-AnimeSharp.pth?download=true", # https://openmodeldb.info/models/4x-AnimeSharp
34
- "4x_IllustrationJaNai_V1_ESRGAN_135k.pth": "https://drive.google.com/uc?export=download&confirm=1&id=1qpioSqBkB_IkSBhEAewSSNFt6qgkBimP", # https://openmodeldb.info/models/4x-IllustrationJaNai-V1-DAT2
 
 
 
 
 
 
 
 
35
  # DATNet
36
- "4xNomos8kDAT.pth": "https://github.com/Phhofm/models/releases/download/4xNomos8kDAT/4xNomos8kDAT.pth", # https://openmodeldb.info/models/4x-Nomos8kDAT
37
- "4x-DWTP-DS-dat2-v3.pth": "https://objectstorage.us-phoenix-1.oraclecloud.com/n/ax6ygfvpvzka/b/open-modeldb-files/o/4x-DWTP-DS-dat2-v3.pth", # https://openmodeldb.info/models/4x-DWTP-DS-dat2-v3
38
- "4x_IllustrationJaNai_V1_DAT2_190k.pth": "https://drive.google.com/uc?export=download&confirm=1&id=1qpioSqBkB_IkSBhEAewSSNFt6qgkBimP", # https://openmodeldb.info/models/4x-IllustrationJaNai-V1-DAT2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39
  # HAT
40
- "4xNomos8kSCHAT-L.pth": "https://github.com/Phhofm/models/releases/download/4xNomos8kSCHAT/4xNomos8kSCHAT-L.pth", # https://openmodeldb.info/models/4x-Nomos8kSCHAT-L
41
- "4xNomos8kSCHAT-S.pth": "https://github.com/Phhofm/models/releases/download/4xNomos8kSCHAT/4xNomos8kSCHAT-S.pth", # https://openmodeldb.info/models/4x-Nomos8kSCHAT-S
42
- "4xNomos8kHAT-L_otf.pth": "https://github.com/Phhofm/models/releases/download/4xNomos8kHAT-L_otf/4xNomos8kHAT-L_otf.pth", # https://openmodeldb.info/models/4x-Nomos8kHAT-L-otf
 
 
 
 
 
 
 
 
 
43
  # RealPLKSR_dysample
44
- "4xHFA2k_ludvae_realplksr_dysample.pth": "https://github.com/Phhofm/models/releases/download/4xHFA2k_ludvae_realplksr_dysample/4xHFA2k_ludvae_realplksr_dysample.pth", # https://openmodeldb.info/models/4x-HFA2k-ludvae-realplksr-dysample
45
- "4xArtFaces_realplksr_dysample.pth": "https://github.com/Phhofm/models/releases/download/4xArtFaces_realplksr_dysample/4xArtFaces_realplksr_dysample.pth", # https://openmodeldb.info/models/4x-ArtFaces-realplksr-dysample
46
- "4x-PBRify_RPLKSRd_V3.pth": "https://github.com/Kim2091/Kim2091-Models/releases/download/4x-PBRify_RPLKSRd_V3/4x-PBRify_RPLKSRd_V3.pth", # https://openmodeldb.info/models/4x-PBRify-RPLKSRd-V3
47
- "4xNomos2_realplksr_dysample.pth": "https://github.com/Phhofm/models/releases/download/4xNomos2_realplksr_dysample/4xNomos2_realplksr_dysample.pth", # https://openmodeldb.info/models/4x-Nomos2-realplksr-dysample
 
 
 
 
 
 
 
 
 
 
 
 
48
  # RealPLKSR
49
- "2x-AnimeSharpV2_RPLKSR_Sharp.pth": "https://github.com/Kim2091/Kim2091-Models/releases/download/2x-AnimeSharpV2_Set/2x-AnimeSharpV2_RPLKSR_Sharp.pth", # https://openmodeldb.info/models/2x-AnimeSharpV2-RPLKSR-Sharp
50
- "2x-AnimeSharpV2_RPLKSR_Soft.pth": "https://github.com/Kim2091/Kim2091-Models/releases/download/2x-AnimeSharpV2_Set/2x-AnimeSharpV2_RPLKSR_Soft.pth", # https://openmodeldb.info/models/2x-AnimeSharpV2-RPLKSR-Soft
51
- "4xPurePhoto-RealPLSKR.pth": "https://github.com/starinspace/StarinspaceUpscale/releases/download/Models/4xPurePhoto-RealPLSKR.pth", # https://openmodeldb.info/models/4x-PurePhoto-RealPLSKR
52
- "2x_Text2HD_v.1-RealPLKSR.pth": "https://github.com/starinspace/StarinspaceUpscale/releases/download/Models/2x_Text2HD_v.1-RealPLKSR.pth", # https://openmodeldb.info/models/2x-Text2HD-v-1
53
- "2xVHS2HD-RealPLKSR.pth": "https://github.com/starinspace/StarinspaceUpscale/releases/download/Models/2xVHS2HD-RealPLKSR.pth", # https://openmodeldb.info/models/2x-VHS2HD
54
- "4xNomosWebPhoto_RealPLKSR.pth": "https://github.com/Phhofm/models/releases/download/4xNomosWebPhoto_RealPLKSR/4xNomosWebPhoto_RealPLKSR.pth", # https://openmodeldb.info/models/4x-NomosWebPhoto-RealPLKSR
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55
  }
56
 
57
- files_to_download = {
58
- "a1.jpg":
59
- "https://thumbs.dreamstime.com/b/tower-bridge-traditional-red-bus-black-white-colors-view-to-tower-bridge-london-black-white-colors-108478942.jpg",
60
- "a2.jpg":
61
- "https://media.istockphoto.com/id/523514029/photo/london-skyline-b-w.jpg?s=612x612&w=0&k=20&c=kJS1BAtfqYeUDaORupj0sBPc1hpzJhBUUqEFfRnHzZ0=",
62
- "a3.jpg":
63
- "https://i.guim.co.uk/img/media/06f614065ed82ca0e917b149a32493c791619854/0_0_3648_2789/master/3648.jpg?width=700&quality=85&auto=format&fit=max&s=05764b507c18a38590090d987c8b6202",
64
- "a4.jpg":
65
- "https://i.pinimg.com/736x/46/96/9e/46969eb94aec2437323464804d27706d--victorian-london-victorian-era.jpg",
66
- }
67
 
68
  def get_model_type(model_name):
69
  # Define model type mappings based on key parts of the model names
@@ -82,19 +230,19 @@ def get_model_type(model_name):
82
  model_type = "RealPLKSR_dysample"
83
  elif "realplksr" in model_name.lower() or "rplksr" in model_name.lower():
84
  model_type = "RealPLKSR"
 
 
 
 
85
  return f"{model_type}, {model_name}"
86
 
87
- typed_realesr_model = {get_model_type(key): value for key, value in realesr_model.items()}
88
-
89
- def download_from_urls(urls, save_dir=None):
90
- for file_name, url in urls.items():
91
- download_from_url(url, file_name, save_dir)
92
 
93
 
94
  class Upscale:
95
- def inference(self, img, face_restoration, realesr, scale: float):
96
  print(img)
97
- print(face_restoration, realesr, scale)
98
  try:
99
  self.scale = scale
100
  self.img_name = os.path.basename(str(img))
@@ -107,55 +255,57 @@ class Upscale:
107
  img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
108
 
109
  h, w = img.shape[0:2]
110
- if h < 300:
111
- img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
112
 
113
  if face_restoration:
114
- download_from_url(face_model[face_restoration], face_restoration, os.path.join("weights", "face"))
115
- if realesr:
116
- realesr_type, realesr = realesr.split(", ", 1)
117
- download_from_url(realesr_model[realesr], realesr, os.path.join("weights", "realesr"))
 
 
 
 
118
 
119
  netscale = 4
120
  loadnet = None
121
  model = None
122
  is_auto_split_upscale = True
123
  half = True if torch.cuda.is_available() else False
124
- if realesr_type:
125
  from basicsr.archs.rrdbnet_arch import RRDBNet
126
  from basicsr.archs.realplksr_arch import realplksr
127
- # background enhancer with RealESRGAN
128
- if realesr_type == "RRDB":
129
- netscale = 2 if "x2" in realesr else 4
130
- num_block = 6 if "6B" in realesr else 23
131
  model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=num_block, num_grow_ch=32, scale=netscale)
132
- elif realesr_type == "SRVGG":
133
  from realesrgan.archs.srvgg_arch import SRVGGNetCompact
134
  netscale = 4
135
- num_conv = 16 if "animevideov3" in realesr else 32
136
  model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=num_conv, upscale=netscale, act_type='prelu')
137
- elif realesr_type == "ESRGAN":
138
  netscale = 4
139
  model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=netscale)
140
  loadnet = {}
141
- loadnet_origin = torch.load(os.path.join("weights", "realesr", realesr), map_location=torch.device('cpu'), weights_only=True)
142
  for key, value in loadnet_origin.items():
143
  new_key = key.replace("model.0", "conv_first").replace("model.1.sub.23.", "conv_body.").replace("model.1.sub", "body") \
144
  .replace(".0.weight", ".weight").replace(".0.bias", ".bias").replace(".RDB1.", ".rdb1.").replace(".RDB2.", ".rdb2.").replace(".RDB3.", ".rdb3.") \
145
  .replace("model.3.", "conv_up1.").replace("model.6.", "conv_up2.").replace("model.8.", "conv_hr.").replace("model.10.", "conv_last.")
146
  loadnet[new_key] = value
147
- elif realesr_type == "DAT":
148
  from basicsr.archs.dat_arch import DAT
149
  half = False
150
  netscale = 4
151
- expansion_factor = 2. if "dat2" in realesr.lower() else 4.
152
  model = DAT(img_size=64, in_chans=3, embed_dim=180, split_size=[8,32], depth=[6,6,6,6,6,6], num_heads=[6,6,6,6,6,6], expansion_factor=expansion_factor, upscale=netscale)
153
  # # Speculate on the parameters.
154
- # loadnet_origin = torch.load(os.path.join("weights", "realesr", realesr), map_location=torch.device('cpu'), weights_only=True)
155
  # inferred_params = self.infer_parameters_from_state_dict_for_dat(loadnet_origin, netscale)
156
  # for param, value in inferred_params.items():
157
  # print(f"{param}: {value}")
158
- elif realesr_type == "HAT":
159
  half = False
160
  netscale = 4
161
  import torch.nn.functional as F
@@ -205,7 +355,7 @@ class Upscale:
205
 
206
  # The parameters are derived from the XPixelGroup project files: HAT-L_SRx4_ImageNet-pretrain.yml and HAT-S_SRx4.yml.
207
  # https://github.com/XPixelGroup/HAT/tree/main/options/test
208
- if "hat-l" in realesr.lower():
209
  window_size = 16
210
  compress_ratio = 3
211
  squeeze_factor = 30
@@ -214,7 +364,7 @@ class Upscale:
214
  num_heads = [6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6]
215
  mlp_ratio = 2
216
  upsampler = "pixelshuffle"
217
- elif "hat-s" in realesr.lower():
218
  window_size = 16
219
  compress_ratio = 24
220
  squeeze_factor = 24
@@ -225,12 +375,12 @@ class Upscale:
225
  upsampler = "pixelshuffle"
226
  model = HATWithAutoPadding(img_size=64, patch_size=1, in_chans=3, embed_dim=embed_dim, depths=depths, num_heads=num_heads, window_size=window_size, compress_ratio=compress_ratio,
227
  squeeze_factor=squeeze_factor, conv_scale=0.01, overlap_ratio=0.5, mlp_ratio=mlp_ratio, upsampler=upsampler, upscale=netscale,)
228
- elif realesr_type == "RealPLKSR_dysample":
229
  netscale = 4
230
- model = realplksr(upscaling_factor=netscale, dysample=True)
231
- elif realesr_type == "RealPLKSR":
232
- half = False if "RealPLSKR" in realesr else half
233
- netscale = 2 if realesr.startswith("2x") else 4
234
  model = realplksr(dim=64, n_blocks=28, kernel_size=17, split_ratio=0.25, upscaling_factor=netscale)
235
 
236
 
@@ -238,8 +388,8 @@ class Upscale:
238
  if loadnet:
239
  self.upsampler = RealESRGANer(scale=netscale, loadnet=loadnet, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
240
  elif model:
241
- self.upsampler = RealESRGANer(scale=netscale, model_path=os.path.join("weights", "realesr", realesr), model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
242
- elif realesr:
243
  self.upsampler = None
244
  import PIL
245
  from image_gen_aux import UpscaleWithModel
@@ -286,25 +436,42 @@ class Upscale:
286
  return cv_image, None
287
 
288
  device = "cuda" if torch.cuda.is_available() else "cpu"
289
- upscaler = UpscaleWithModel.from_pretrained(os.path.join("weights", "realesr", realesr)).to(device)
290
  upscaler.__class__ = UpscaleWithModel_Gfpgan
291
  self.upsampler = upscaler
292
  self.face_enhancer = None
293
 
 
294
  if face_restoration:
 
295
  from gfpgan.utils import GFPGANer
 
 
 
 
296
  if face_restoration and face_restoration.startswith("GFPGANv1."):
297
- self.face_enhancer = GFPGANer(model_path=os.path.join("weights", "face", face_restoration), upscale=self.scale, arch="clean", channel_multiplier=2, bg_upsampler=self.upsampler)
 
298
  elif face_restoration and face_restoration.startswith("RestoreFormer"):
299
  arch = "RestoreFormer++" if face_restoration.startswith("RestoreFormer++") else "RestoreFormer"
300
- self.face_enhancer = GFPGANer(model_path=os.path.join("weights", "face", face_restoration), upscale=self.scale, arch=arch, channel_multiplier=2, bg_upsampler=self.upsampler)
301
  elif face_restoration == 'CodeFormer.pth':
302
- self.face_enhancer = GFPGANer(
303
- model_path='weights/CodeFormer.pth', upscale=self.scale, arch='CodeFormer', channel_multiplier=2, bg_upsampler=self.upsampler)
304
-
 
 
 
 
 
 
 
 
 
305
 
306
  files = []
307
- outputs = []
 
 
308
  try:
309
  bg_upsample_img = None
310
  if self.upsampler and self.upsampler.enhance:
@@ -317,22 +484,19 @@ class Upscale:
317
  if cropped_faces and restored_aligned:
318
  for idx, (cropped_face, restored_face) in enumerate(zip(cropped_faces, restored_aligned)):
319
  # save cropped face
320
- save_crop_path = f"output/{self.basename}{idx:02d}_cropped_faces.png"
321
  self.imwriteUTF8(save_crop_path, cropped_face)
322
  # save restored face
323
- save_restore_path = f"output/{self.basename}{idx:02d}_restored_faces.png"
324
  self.imwriteUTF8(save_restore_path, restored_face)
325
  # save comparison image
326
- save_cmp_path = f"output/{self.basename}{idx:02d}_cmp.png"
327
  cmp_img = np.concatenate((cropped_face, restored_face), axis=1)
328
  self.imwriteUTF8(save_cmp_path, cmp_img)
329
 
330
  files.append(save_crop_path)
331
  files.append(save_restore_path)
332
  files.append(save_cmp_path)
333
- outputs.append(cv2.cvtColor(cropped_face, cv2.COLOR_BGR2RGB))
334
- outputs.append(cv2.cvtColor(restored_face, cv2.COLOR_BGR2RGB))
335
- outputs.append(cv2.cvtColor(cmp_img, cv2.COLOR_BGR2RGB))
336
 
337
  restored_img = bg_upsample_img
338
  except RuntimeError as error:
@@ -348,13 +512,12 @@ class Upscale:
348
 
349
  if not self.extension:
350
  self.extension = ".png" if self.img_mode == "RGBA" else ".jpg" # RGBA images should be saved in png format
351
- save_path = f"output/{self.basename}{self.extension}"
352
  self.imwriteUTF8(save_path, restored_img)
353
 
354
  restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
355
  files.append(save_path)
356
- outputs.append(restored_img)
357
- return outputs, files
358
  except Exception as error:
359
  print(traceback.format_exc())
360
  print("global exception", error)
@@ -451,42 +614,87 @@ def main():
451
  # Ensure the target directory exists
452
  os.makedirs('output', exist_ok=True)
453
 
454
- # Iterate through each file
455
- download_from_urls(files_to_download, ".")
456
-
457
- title = "Image Upscaling & Restoration(esp. Face) using GFPGAN Algorithm"
458
- description = r"""Gradio demo for <a href='https://github.com/TencentARC/GFPGAN' target='_blank'><b>GFPGAN: Towards Real-World Blind Face Restoration and Upscalling of the image with a Generative Facial Prior</b></a>.<br>
459
- Practically the algorithm is used to restore your **old photos** or improve **AI-generated faces**.<br>
 
 
460
  To use it, simply just upload the concerned image.<br>
461
  """
462
  article = r"""
463
  [![download](https://img.shields.io/github/downloads/TencentARC/GFPGAN/total.svg)](https://github.com/TencentARC/GFPGAN/releases)
464
  [![GitHub Stars](https://img.shields.io/github/stars/TencentARC/GFPGAN?style=social)](https://github.com/TencentARC/GFPGAN)
465
  [![arXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/abs/2101.04061)
466
- <center><img src='https://visitor-badge.glitch.me/badge?page_id=dj_face_restoration_GFPGAN' alt='visitor badge'></center>
467
  """
468
 
469
  upscale = Upscale()
470
 
471
- demo = gr.Interface(
472
- upscale.inference, [
473
- gr.Image(type="filepath", label="Input", format="png"),
474
- gr.Dropdown(list(face_model.keys())+[None], type="value", value='GFPGANv1.4.pth', label='Face Restoration version', info="Face Restoration and RealESR can be freely combined in different ways, or one can be set to \"None\" to use only the other model. Face Restoration is primarily used for face restoration in real-life images, while RealESR serves as a background restoration model."),
475
- gr.Dropdown(list(typed_realesr_model.keys())+[None], type="value", value='SRVGG, realesr-general-x4v3.pth', label='RealESR version'),
476
- gr.Number(label="Rescaling factor", value=4),
477
- ], [
478
- gr.Gallery(type="numpy", label="Output (The whole image)", format="png"),
479
- gr.File(label="Download the output image")
480
- ],
481
- title=title,
482
- description=description,
483
- article=article,
484
- examples=[["a1.jpg", "GFPGANv1.4.pth", "SRVGG, realesr-general-x4v3.pth", 2],
485
- ["a2.jpg", "GFPGANv1.4.pth", "SRVGG, realesr-general-x4v3.pth", 2],
486
- ["a3.jpg", "GFPGANv1.4.pth", "SRVGG, realesr-general-x4v3.pth", 2],
487
- ["a4.jpg", "GFPGANv1.4.pth", "SRVGG, realesr-general-x4v3.pth", 2]])
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
488
 
489
- demo.queue(default_concurrency_limit=4)
490
  demo.launch(inbrowser=True)
491
 
492
 
 
1
  import os
2
  import gc
3
  import cv2
 
4
  import numpy as np
5
  import gradio as gr
6
  import torch
 
10
 
11
 
12
  # Define URLs and their corresponding local storage paths
13
+ face_models = {
14
+ "GFPGANv1.4.pth" : ["https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth",
15
+ "https://github.com/TencentARC/GFPGAN/",
16
+ """GFPGAN: Towards Real-World Blind Face Restoration and Upscalling of the image with a Generative Facial Prior.
17
+ GFPGAN aims at developing a Practical Algorithm for Real-world Face Restoration.
18
+ It leverages rich and diverse priors encapsulated in a pretrained face GAN (e.g., StyleGAN2) for blind face restoration."""],
19
+
20
+ "RestoreFormer++.ckpt": ["https://github.com/wzhouxiff/RestoreFormerPlusPlus/releases/download/v1.0.0/RestoreFormer++.ckpt",
21
+ "https://github.com/wzhouxiff/RestoreFormerPlusPlus",
22
+ """RestoreFormer++: Towards Real-World Blind Face Restoration from Undegraded Key-Value Pairs.
23
+ RestoreFormer++ is an extension of RestoreFormer. It proposes to restore a degraded face image with both fidelity and \
24
+ realness by using the powerful fully-spacial attention mechanisms to model the abundant contextual information in the face and \
25
+ its interplay with reconstruction-oriented high-quality priors."""],
26
+
27
+ "CodeFormer.pth" : ["https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth",
28
+ "https://github.com/sczhou/CodeFormer",
29
+ """CodeFormer: Towards Robust Blind Face Restoration with Codebook Lookup Transformer (NeurIPS 2022).
30
+ CodeFormer is a Transformer-based model designed to tackle the challenging problem of blind face restoration, where inputs are often severely degraded.
31
+ By framing face restoration as a code prediction task, this approach ensures both improved mapping from degraded inputs to outputs and the generation of visually rich, high-quality faces.
32
+ """],
33
+
34
+ "GPEN-BFR-512.pth" : ["https://huggingface.co/akhaliq/GPEN-BFR-512/resolve/main/GPEN-BFR-512.pth",
35
+ "https://github.com/yangxy/GPEN",
36
+ """GPEN: GAN Prior Embedded Network for Blind Face Restoration in the Wild.
37
+ GPEN addresses blind face restoration (BFR) by embedding a GAN into a U-shaped DNN, combining GAN’s ability to generate high-quality images with DNN’s feature extraction.
38
+ This design reconstructs global structure, fine details, and backgrounds from degraded inputs.
39
+ Simple yet effective, GPEN outperforms state-of-the-art methods, delivering realistic results even for severely degraded images."""],
40
+
41
+ "GPEN-BFR-1024.pt" : ["https://www.modelscope.cn/models/iic/cv_gpen_image-portrait-enhancement-hires/resolve/master/pytorch_model.pt",
42
+ "https://www.modelscope.cn/models/iic/cv_gpen_image-portrait-enhancement-hires/files",
43
+ """The same as GPEN but for 1024 resolution."""],
44
+
45
+ "GPEN-BFR-2048.pt" : ["https://www.modelscope.cn/models/iic/cv_gpen_image-portrait-enhancement-hires/resolve/master/pytorch_model-2048.pt",
46
+ "https://www.modelscope.cn/models/iic/cv_gpen_image-portrait-enhancement-hires/files",
47
+ """The same as GPEN but for 2048 resolution."""],
48
+
49
  # legacy model
50
+ "GFPGANv1.3.pth" : ["https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth",
51
+ "https://github.com/TencentARC/GFPGAN/", "The same as GFPGAN but legacy model"],
52
+ "GFPGANv1.2.pth" : ["https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.2.pth",
53
+ "https://github.com/TencentARC/GFPGAN/", "The same as GFPGAN but legacy model"],
54
+ "RestoreFormer.ckpt": ["https://github.com/wzhouxiff/RestoreFormerPlusPlus/releases/download/v1.0.0/RestoreFormer.ckpt",
55
+ "https://github.com/wzhouxiff/RestoreFormerPlusPlus", "The same as RestoreFormer++ but legacy model"],
56
  }
57
+ upscale_models = {
58
  # SRVGGNet
59
+ "realesr-general-x4v3.pth": ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
60
+ "https://github.com/xinntao/Real-ESRGAN/releases/tag/v0.3.0",
61
+ """add realesr-general-x4v3 and realesr-general-wdn-x4v3. They are very tiny models for general scenes, and they may more robust. But as they are tiny models, their performance may be limited."""],
62
+
63
+ "realesr-animevideov3.pth": ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth",
64
+ "https://github.com/xinntao/Real-ESRGAN/releases/tag/v0.2.5.0",
65
+ """update the RealESRGAN AnimeVideo-v3 model, which can achieve better results with a faster inference speed."""],
66
+
67
  # RRDBNet
68
+ "RealESRGAN_x4plus_anime_6B.pth": ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
69
+ "https://github.com/xinntao/Real-ESRGAN/releases/tag/v0.2.2.4",
70
+ """We add RealESRGAN_x4plus_anime_6B.pth, which is optimized for anime images with much smaller model size. More details and comparisons with waifu2x are in anime_model.md"""],
71
+
72
+ "RealESRGAN_x2plus.pth" : ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
73
+ "https://github.com/xinntao/Real-ESRGAN/releases/tag/v0.2.1",
74
+ """Add RealESRGAN_x2plus.pth model"""],
75
+
76
+ "RealESRNet_x4plus.pth" : ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth",
77
+ "https://github.com/xinntao/Real-ESRGAN/releases/tag/v0.1.1",
78
+ """This release is mainly for storing pre-trained models and executable files."""],
79
+
80
+ "RealESRGAN_x4plus.pth" : ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
81
+ "https://github.com/xinntao/Real-ESRGAN/releases/tag/v0.1.0",
82
+ """This release is mainly for storing pre-trained models and executable files."""],
83
+
84
  # ESRGAN(oldRRDB)
85
+ "4x-AnimeSharp.pth": ["https://huggingface.co/utnah/esrgan/resolve/main/4x-AnimeSharp.pth?download=true",
86
+ "https://openmodeldb.info/models/4x-AnimeSharp",
87
+ """Interpolation between 4x-UltraSharp and 4x-TextSharp-v0.5. Works amazingly on anime. It also upscales text, but it's far better with anime content."""],
88
+
89
+ "4x_IllustrationJaNai_V1_ESRGAN_135k.pth": ["https://drive.google.com/uc?export=download&confirm=1&id=1qpioSqBkB_IkSBhEAewSSNFt6qgkBimP",
90
+ "https://openmodeldb.info/models/4x-IllustrationJaNai-V1-DAT2",
91
+ """Purpose: Illustrations, digital art, manga covers
92
+ Model for color images including manga covers and color illustrations, digital art, visual novel art, artbooks, and more.
93
+ DAT2 version is the highest quality version but also the slowest. See the ESRGAN version for faster performance."""],
94
+
95
  # DATNet
96
+ "4xNomos8kDAT.pth" : ["https://github.com/Phhofm/models/releases/download/4xNomos8kDAT/4xNomos8kDAT.pth",
97
+ "https://openmodeldb.info/models/4x-Nomos8kDAT",
98
+ """A 4x photo upscaler with otf jpg compression, blur and resize, trained on musl's Nomos8k_sfw dataset for realisic sr, this time based on the DAT arch, as a finetune on the official 4x DAT model."""],
99
+
100
+ "4x-DWTP-DS-dat2-v3.pth" : ["https://objectstorage.us-phoenix-1.oraclecloud.com/n/ax6ygfvpvzka/b/open-modeldb-files/o/4x-DWTP-DS-dat2-v3.pth",
101
+ "https://openmodeldb.info/models/4x-DWTP-DS-dat2-v3",
102
+ """DAT descreenton model, designed to reduce discrepancies on tiles due to too much loss of the first version, while getting rid of the removal of paper texture"""],
103
+
104
+ "4xBHI_dat2_real.pth" : ["https://github.com/Phhofm/models/releases/download/4xBHI_dat2_real/4xBHI_dat2_real.pth",
105
+ "https://github.com/Phhofm/models/releases/tag/4xBHI_dat2_real",
106
+ """Purpose: 4x upscaling images. Handles realistic noise, some realistic blur, and webp and jpg (re)compression.
107
+ Description: 4x dat2 upscaling model for web and realistic images. It handles realistic noise, some realistic blur, and webp and jpg (re)compression. Trained on my BHI dataset (390'035 training tiles) with degraded LR subset."""],
108
+
109
+ "4xBHI_dat2_otf.pth" : ["https://github.com/Phhofm/models/releases/download/4xBHI_dat2_otf/4xBHI_dat2_otf.pth",
110
+ "https://github.com/Phhofm/models/releases/tag/4xBHI_dat2_otf",
111
+ """Purpose: 4x upscaling images, handles noise and jpg compression
112
+ Description: 4x dat2 upscaling model, trained with the real-esrgan otf pipeline on my bhi dataset. Handles noise and compression."""],
113
+
114
+ "4xBHI_dat2_multiblur.pth" : ["https://github.com/Phhofm/models/releases/download/4xBHI_dat2_multiblurjpg/4xBHI_dat2_multiblur.pth",
115
+ "https://github.com/Phhofm/models/releases/tag/4xBHI_dat2_multiblurjpg",
116
+ """Purpose: 4x upscaling images, handles jpg compression
117
+ Description: 4x dat2 upscaling model, trained with down_up,linear, cubic_mitchell, lanczos, gauss and box scaling algos, some average, gaussian and anisotropic blurs and jpg compression. Trained on my BHI sisr dataset."""],
118
+
119
+ "4xBHI_dat2_multiblurjpg.pth" : ["https://github.com/Phhofm/models/releases/download/4xBHI_dat2_multiblurjpg/4xBHI_dat2_multiblurjpg.pth",
120
+ "https://github.com/Phhofm/models/releases/tag/4xBHI_dat2_multiblurjpg",
121
+ """Purpose: 4x upscaling images, handles jpg compression
122
+ Description: 4x dat2 upscaling model, trained with down_up,linear, cubic_mitchell, lanczos, gauss and box scaling algos, some average, gaussian and anisotropic blurs and jpg compression. Trained on my BHI sisr dataset."""],
123
+
124
+ "4x_IllustrationJaNai_V1_DAT2_190k.pth": ["https://drive.google.com/uc?export=download&confirm=1&id=1qpioSqBkB_IkSBhEAewSSNFt6qgkBimP",
125
+ "https://openmodeldb.info/models/4x-IllustrationJaNai-V1-DAT2",
126
+ """Purpose: Illustrations, digital art, manga covers
127
+ Model for color images including manga covers and color illustrations, digital art, visual novel art, artbooks, and more.
128
+ DAT2 version is the highest quality version but also the slowest. See the ESRGAN version for faster performance."""],
129
+
130
  # HAT
131
+ "4xNomos8kSCHAT-L.pth" : ["https://github.com/Phhofm/models/releases/download/4xNomos8kSCHAT/4xNomos8kSCHAT-L.pth",
132
+ "https://openmodeldb.info/models/4x-Nomos8kSCHAT-L",
133
+ """4x photo upscaler with otf jpg compression and blur, trained on musl's Nomos8k_sfw dataset for realisic sr. Since this is a big model, upscaling might take a while."""],
134
+
135
+ "4xNomos8kSCHAT-S.pth" : ["https://github.com/Phhofm/models/releases/download/4xNomos8kSCHAT/4xNomos8kSCHAT-S.pth",
136
+ "https://openmodeldb.info/models/4x-Nomos8kSCHAT-S",
137
+ """4x photo upscaler with otf jpg compression and blur, trained on musl's Nomos8k_sfw dataset for realisic sr. HAT-S version/model."""],
138
+
139
+ "4xNomos8kHAT-L_otf.pth": ["https://github.com/Phhofm/models/releases/download/4xNomos8kHAT-L_otf/4xNomos8kHAT-L_otf.pth",
140
+ "https://openmodeldb.info/models/4x-Nomos8kHAT-L-otf",
141
+ """4x photo upscaler trained with otf"""],
142
+
143
  # RealPLKSR_dysample
144
+ "4xHFA2k_ludvae_realplksr_dysample.pth": ["https://github.com/Phhofm/models/releases/download/4xHFA2k_ludvae_realplksr_dysample/4xHFA2k_ludvae_realplksr_dysample.pth",
145
+ "https://openmodeldb.info/models/4x-HFA2k-ludvae-realplksr-dysample",
146
+ """A Dysample RealPLKSR 4x upscaling model for anime single-image resolution."""],
147
+
148
+ "4xArtFaces_realplksr_dysample.pth" : ["https://github.com/Phhofm/models/releases/download/4xArtFaces_realplksr_dysample/4xArtFaces_realplksr_dysample.pth",
149
+ "https://openmodeldb.info/models/4x-ArtFaces-realplksr-dysample",
150
+ """A Dysample RealPLKSR 4x upscaling model for art / painted faces."""],
151
+
152
+ "4x-PBRify_RPLKSRd_V3.pth" : ["https://github.com/Kim2091/Kim2091-Models/releases/download/4x-PBRify_RPLKSRd_V3/4x-PBRify_RPLKSRd_V3.pth", "https://openmodeldb.info/models/4x-PBRify-RPLKSRd-V3",
153
+ """This model is roughly 8x faster than the current DAT2 model, while being higher quality. It produces far more natural detail, resolves lines and edges more smoothly, and cleans up compression artifacts better."""],
154
+
155
+ "4xNomos2_realplksr_dysample.pth" : ["https://github.com/Phhofm/models/releases/download/4xNomos2_realplksr_dysample/4xNomos2_realplksr_dysample.pth",
156
+ "https://openmodeldb.info/models/4x-Nomos2-realplksr-dysample",
157
+ """Description: A Dysample RealPLKSR 4x upscaling model that was trained with / handles jpg compression down to 70 on the Nomosv2 dataset, preserves DoF.
158
+ This model affects / saturate colors, which can be counteracted a bit by using wavelet color fix, as used in these examples."""],
159
+
160
  # RealPLKSR
161
+ "2x-AnimeSharpV2_RPLKSR_Sharp.pth": ["https://github.com/Kim2091/Kim2091-Models/releases/download/2x-AnimeSharpV2_Set/2x-AnimeSharpV2_RPLKSR_Sharp.pth",
162
+ "https://openmodeldb.info/models/2x-AnimeSharpV2-RPLKSR-Sharp",
163
+ """Kim2091: This is my first anime model in years. Hopefully you guys can find a good use-case for it.
164
+ RealPLKSR (Higher quality, slower) Sharp: For heavily degraded sources. Sharp models have issues depth of field but are best at removing artifacts
165
+ """],
166
+
167
+ "2x-AnimeSharpV2_RPLKSR_Soft.pth" : ["https://github.com/Kim2091/Kim2091-Models/releases/download/2x-AnimeSharpV2_Set/2x-AnimeSharpV2_RPLKSR_Soft.pth",
168
+ "https://openmodeldb.info/models/2x-AnimeSharpV2-RPLKSR-Soft",
169
+ """Kim2091: This is my first anime model in years. Hopefully you guys can find a good use-case for it.
170
+ RealPLKSR (Higher quality, slower) Soft: For cleaner sources. Soft models preserve depth of field but may not remove other artifacts as well"""],
171
+
172
+ "4xPurePhoto-RealPLSKR.pth" : ["https://github.com/starinspace/StarinspaceUpscale/releases/download/Models/4xPurePhoto-RealPLSKR.pth",
173
+ "https://openmodeldb.info/models/4x-PurePhoto-RealPLSKR",
174
+ """Skilled in working with cats, hair, parties, and creating clear images.
175
+ Also proficient in resizing photos and enlarging large, sharp images.
176
+ Can effectively improve images from small sizes as well (300px at smallest on one side, depending on the subject).
177
+ Experienced in experimenting with techniques like upscaling with this model twice and \
178
+ then reducing it by 50% to enhance details, especially in features like hair or animals."""],
179
+
180
+ "2x_Text2HD_v.1-RealPLKSR.pth" : ["https://github.com/starinspace/StarinspaceUpscale/releases/download/Models/2x_Text2HD_v.1-RealPLKSR.pth",
181
+ "https://openmodeldb.info/models/2x-Text2HD-v-1",
182
+ """Purpose: Upscale text in very low quality to normal quality.
183
+ The upscale model is specifically designed to enhance lower-quality text images, \
184
+ improving their clarity and readability by upscaling them by 2x.
185
+ It excels at processing moderately sized text, effectively transforming it into high-quality, legible scans.
186
+ However, the model may encounter challenges when dealing with very small text, \
187
+ as its performance is optimized for text of a certain minimum size. For best results, \
188
+ input images should contain text that is not excessively small."""],
189
+
190
+ "2xVHS2HD-RealPLKSR.pth" : ["https://github.com/starinspace/StarinspaceUpscale/releases/download/Models/2xVHS2HD-RealPLKSR.pth",
191
+ "https://openmodeldb.info/models/2x-VHS2HD",
192
+ """An advanced VHS recording model designed to enhance video quality by reducing artifacts such as haloing, ghosting, and noise patterns.
193
+ Optimized primarily for PAL resolution (NTSC might work good as well)."""],
194
+
195
+ "4xNomosWebPhoto_RealPLKSR.pth" : ["https://github.com/Phhofm/models/releases/download/4xNomosWebPhoto_RealPLKSR/4xNomosWebPhoto_RealPLKSR.pth",
196
+ "https://openmodeldb.info/models/4x-NomosWebPhoto-RealPLKSR",
197
+ """4x RealPLKSR model for photography, trained with realistic noise, lens blur, jpg and webp re-compression."""],
198
+
199
+ # "4xNomos2_hq_drct-l.pth" : ["https://github.com/Phhofm/models/releases/download/4xNomos2_hq_drct-l/4xNomos2_hq_drct-l.pth",
200
+ # "https://github.com/Phhofm/models/releases/tag/4xNomos2_hq_drct-l",
201
+ # """An drct-l 4x upscaling model, similiar to the 4xNomos2_hq_atd, 4xNomos2_hq_dat2 and 4xNomos2_hq_mosr models, trained and for usage on non-degraded input to give good quality output.
202
+ # """],
203
+
204
+ # "4xNomos2_hq_atd.pth" : ["https://github.com/Phhofm/models/releases/download/4xNomos2_hq_atd/4xNomos2_hq_atd.pth",
205
+ # "https://github.com/Phhofm/models/releases/tag/4xNomos2_hq_atd",
206
+ # """An atd 4x upscaling model, similiar to the 4xNomos2_hq_dat2 or 4xNomos2_hq_mosr models, trained and for usage on non-degraded input to give good quality output.
207
+ # """]
208
  }
209
 
210
+ example_list = ["images/a01.jpg", "images/a02.jpg", "images/a03.jpg", "images/a04.jpg", "images/bus.jpg", "images/zidane.jpg",
211
+ "images/b01.jpg", "images/b02.jpg", "images/b03.jpg", "images/b04.jpg", "images/b05.jpg", "images/b06.jpg",
212
+ "images/b07.jpg", "images/b08.jpg", "images/b09.jpg", "images/b10.jpg", "images/b11.jpg", "images/c01.jpg",
213
+ "images/c02.jpg", "images/c03.jpg", "images/c04.jpg", "images/c05.jpg", "images/c06.jpg", "images/c07.jpg",
214
+ "images/c08.jpg", "images/c09.jpg", "images/c10.jpg"]
 
 
 
 
 
215
 
216
  def get_model_type(model_name):
217
  # Define model type mappings based on key parts of the model names
 
230
  model_type = "RealPLKSR_dysample"
231
  elif "realplksr" in model_name.lower() or "rplksr" in model_name.lower():
232
  model_type = "RealPLKSR"
233
+ elif "drct-l" in model_name.lower():
234
+ model_type = "DRCT-L"
235
+ elif "atd" in model_name.lower():
236
+ model_type = "ATD"
237
  return f"{model_type}, {model_name}"
238
 
239
+ typed_upscale_models = {get_model_type(key): value[0] for key, value in upscale_models.items()}
 
 
 
 
240
 
241
 
242
  class Upscale:
243
+ def inference(self, img, face_restoration, upscale_model, scale: float, face_detection, outputWithModelName: bool):
244
  print(img)
245
+ print(face_restoration, upscale_model, scale)
246
  try:
247
  self.scale = scale
248
  self.img_name = os.path.basename(str(img))
 
255
  img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
256
 
257
  h, w = img.shape[0:2]
 
 
258
 
259
  if face_restoration:
260
+ download_from_url(face_models[face_restoration][0], face_restoration, os.path.join("weights", "face"))
261
+
262
+ modelInUse = ""
263
+ upscale_type = None
264
+ if upscale_model:
265
+ upscale_type, upscale_model = upscale_model.split(", ", 1)
266
+ download_from_url(upscale_models[upscale_model][0], upscale_model, os.path.join("weights", "upscale"))
267
+ modelInUse = f"_{os.path.splitext(upscale_model)[0]}"
268
 
269
  netscale = 4
270
  loadnet = None
271
  model = None
272
  is_auto_split_upscale = True
273
  half = True if torch.cuda.is_available() else False
274
+ if upscale_type:
275
  from basicsr.archs.rrdbnet_arch import RRDBNet
276
  from basicsr.archs.realplksr_arch import realplksr
277
+ # background enhancer with upscale model
278
+ if upscale_type == "RRDB":
279
+ netscale = 2 if "x2" in upscale_model else 4
280
+ num_block = 6 if "6B" in upscale_model else 23
281
  model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=num_block, num_grow_ch=32, scale=netscale)
282
+ elif upscale_type == "SRVGG":
283
  from realesrgan.archs.srvgg_arch import SRVGGNetCompact
284
  netscale = 4
285
+ num_conv = 16 if "animevideov3" in upscale_model else 32
286
  model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=num_conv, upscale=netscale, act_type='prelu')
287
+ elif upscale_type == "ESRGAN":
288
  netscale = 4
289
  model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=netscale)
290
  loadnet = {}
291
+ loadnet_origin = torch.load(os.path.join("weights", "upscale", upscale_model), map_location=torch.device('cpu'), weights_only=True)
292
  for key, value in loadnet_origin.items():
293
  new_key = key.replace("model.0", "conv_first").replace("model.1.sub.23.", "conv_body.").replace("model.1.sub", "body") \
294
  .replace(".0.weight", ".weight").replace(".0.bias", ".bias").replace(".RDB1.", ".rdb1.").replace(".RDB2.", ".rdb2.").replace(".RDB3.", ".rdb3.") \
295
  .replace("model.3.", "conv_up1.").replace("model.6.", "conv_up2.").replace("model.8.", "conv_hr.").replace("model.10.", "conv_last.")
296
  loadnet[new_key] = value
297
+ elif upscale_type == "DAT":
298
  from basicsr.archs.dat_arch import DAT
299
  half = False
300
  netscale = 4
301
+ expansion_factor = 2. if "dat2" in upscale_model.lower() else 4.
302
  model = DAT(img_size=64, in_chans=3, embed_dim=180, split_size=[8,32], depth=[6,6,6,6,6,6], num_heads=[6,6,6,6,6,6], expansion_factor=expansion_factor, upscale=netscale)
303
  # # Speculate on the parameters.
304
+ # loadnet_origin = torch.load(os.path.join("weights", "upscale", upscale_model), map_location=torch.device('cpu'), weights_only=True)
305
  # inferred_params = self.infer_parameters_from_state_dict_for_dat(loadnet_origin, netscale)
306
  # for param, value in inferred_params.items():
307
  # print(f"{param}: {value}")
308
+ elif upscale_type == "HAT":
309
  half = False
310
  netscale = 4
311
  import torch.nn.functional as F
 
355
 
356
  # The parameters are derived from the XPixelGroup project files: HAT-L_SRx4_ImageNet-pretrain.yml and HAT-S_SRx4.yml.
357
  # https://github.com/XPixelGroup/HAT/tree/main/options/test
358
+ if "hat-l" in upscale_model.lower():
359
  window_size = 16
360
  compress_ratio = 3
361
  squeeze_factor = 30
 
364
  num_heads = [6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6]
365
  mlp_ratio = 2
366
  upsampler = "pixelshuffle"
367
+ elif "hat-s" in upscale_model.lower():
368
  window_size = 16
369
  compress_ratio = 24
370
  squeeze_factor = 24
 
375
  upsampler = "pixelshuffle"
376
  model = HATWithAutoPadding(img_size=64, patch_size=1, in_chans=3, embed_dim=embed_dim, depths=depths, num_heads=num_heads, window_size=window_size, compress_ratio=compress_ratio,
377
  squeeze_factor=squeeze_factor, conv_scale=0.01, overlap_ratio=0.5, mlp_ratio=mlp_ratio, upsampler=upsampler, upscale=netscale,)
378
+ elif upscale_type == "RealPLKSR_dysample":
379
  netscale = 4
380
+ model = realplksr(dim=64, n_blocks=28, kernel_size=17, split_ratio=0.25, upscaling_factor=netscale, dysample=True)
381
+ elif upscale_type == "RealPLKSR":
382
+ half = False if "RealPLSKR" in upscale_model else half
383
+ netscale = 2 if upscale_model.startswith("2x") else 4
384
  model = realplksr(dim=64, n_blocks=28, kernel_size=17, split_ratio=0.25, upscaling_factor=netscale)
385
 
386
 
 
388
  if loadnet:
389
  self.upsampler = RealESRGANer(scale=netscale, loadnet=loadnet, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
390
  elif model:
391
+ self.upsampler = RealESRGANer(scale=netscale, model_path=os.path.join("weights", "upscale", upscale_model), model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
392
+ elif upscale_model:
393
  self.upsampler = None
394
  import PIL
395
  from image_gen_aux import UpscaleWithModel
 
436
  return cv_image, None
437
 
438
  device = "cuda" if torch.cuda.is_available() else "cpu"
439
+ upscaler = UpscaleWithModel.from_pretrained(os.path.join("weights", "upscale", upscale_model)).to(device)
440
  upscaler.__class__ = UpscaleWithModel_Gfpgan
441
  self.upsampler = upscaler
442
  self.face_enhancer = None
443
 
444
+ resolution = 512
445
  if face_restoration:
446
+ modelInUse = f"_{os.path.splitext(face_restoration)[0]}" + modelInUse
447
  from gfpgan.utils import GFPGANer
448
+ model_rootpath = os.path.join("weights", "face")
449
+ model_path = os.path.join(model_rootpath, face_restoration)
450
+ channel_multiplier = None
451
+
452
  if face_restoration and face_restoration.startswith("GFPGANv1."):
453
+ arch = "clean"
454
+ channel_multiplier = 2
455
  elif face_restoration and face_restoration.startswith("RestoreFormer"):
456
  arch = "RestoreFormer++" if face_restoration.startswith("RestoreFormer++") else "RestoreFormer"
 
457
  elif face_restoration == 'CodeFormer.pth':
458
+ arch = "CodeFormer"
459
+ elif face_restoration.startswith("GPEN-BFR-"):
460
+ arch = "GPEN"
461
+ channel_multiplier = 2
462
+ if "1024" in face_restoration:
463
+ arch = "GPEN-1024"
464
+ resolution = 1024
465
+ elif "2048" in face_restoration:
466
+ arch = "GPEN-2048"
467
+ resolution = 2048
468
+
469
+ self.face_enhancer = GFPGANer(model_path=model_path, upscale=self.scale, arch=arch, channel_multiplier=channel_multiplier, bg_upsampler=self.upsampler, model_rootpath=model_rootpath, det_model=face_detection, resolution=resolution)
470
 
471
  files = []
472
+ if not outputWithModelName:
473
+ modelInUse = ""
474
+
475
  try:
476
  bg_upsample_img = None
477
  if self.upsampler and self.upsampler.enhance:
 
484
  if cropped_faces and restored_aligned:
485
  for idx, (cropped_face, restored_face) in enumerate(zip(cropped_faces, restored_aligned)):
486
  # save cropped face
487
+ save_crop_path = f"output/{self.basename}{idx:02d}_cropped_faces{modelInUse}.png"
488
  self.imwriteUTF8(save_crop_path, cropped_face)
489
  # save restored face
490
+ save_restore_path = f"output/{self.basename}{idx:02d}_restored_faces{modelInUse}.png"
491
  self.imwriteUTF8(save_restore_path, restored_face)
492
  # save comparison image
493
+ save_cmp_path = f"output/{self.basename}{idx:02d}_cmp{modelInUse}.png"
494
  cmp_img = np.concatenate((cropped_face, restored_face), axis=1)
495
  self.imwriteUTF8(save_cmp_path, cmp_img)
496
 
497
  files.append(save_crop_path)
498
  files.append(save_restore_path)
499
  files.append(save_cmp_path)
 
 
 
500
 
501
  restored_img = bg_upsample_img
502
  except RuntimeError as error:
 
512
 
513
  if not self.extension:
514
  self.extension = ".png" if self.img_mode == "RGBA" else ".jpg" # RGBA images should be saved in png format
515
+ save_path = f"output/{self.basename}{modelInUse}{self.extension}"
516
  self.imwriteUTF8(save_path, restored_img)
517
 
518
  restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
519
  files.append(save_path)
520
+ return files, files
 
521
  except Exception as error:
522
  print(traceback.format_exc())
523
  print("global exception", error)
 
614
  # Ensure the target directory exists
615
  os.makedirs('output', exist_ok=True)
616
 
617
+ title = "Image Upscaling & Restoration using GFPGAN / RestoreFormerPlusPlus / CodeFormer / GPEN Algorithm"
618
+ description = r"""
619
+ <a href='https://github.com/TencentARC/GFPGAN' target='_blank'><b>GFPGAN: Towards Real-World Blind Face Restoration and Upscalling of the image with a Generative Facial Prior</b></a>. <br>
620
+ <a href='https://github.com/wzhouxiff/RestoreFormerPlusPlus' target='_blank'><b>RestoreFormer++: Towards Real-World Blind Face Restoration from Undegraded Key-Value Pairs</b></a>. <br>
621
+ <a href='https://github.com/sczhou/CodeFormer' target='_blank'><b>CodeFormer: Towards Robust Blind Face Restoration with Codebook Lookup Transformer (NeurIPS 2022)</b></a>. <br>
622
+ <a href='https://github.com/yangxy/GPEN' target='_blank'><b>GPEN: GAN Prior Embedded Network for Blind Face Restoration in the Wild</b></a>. <br>
623
+ <br>
624
+ Practically, the aforementioned algorithm is used to restore your **old photos** or improve **AI-generated faces**.<br>
625
  To use it, simply just upload the concerned image.<br>
626
  """
627
  article = r"""
628
  [![download](https://img.shields.io/github/downloads/TencentARC/GFPGAN/total.svg)](https://github.com/TencentARC/GFPGAN/releases)
629
  [![GitHub Stars](https://img.shields.io/github/stars/TencentARC/GFPGAN?style=social)](https://github.com/TencentARC/GFPGAN)
630
  [![arXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/abs/2101.04061)
 
631
  """
632
 
633
  upscale = Upscale()
634
 
635
+ with gr.Blocks(title = title) as demo:
636
+ gr.Markdown(value=f"<h1 style=\"text-align:center;\">{title}</h1><br>{description}")
637
+ with gr.Row():
638
+ with gr.Column(variant="panel"):
639
+ input_image = gr.Image(type="filepath", label="Input", format="png")
640
+ face_model = gr.Dropdown(list(face_models.keys())+[None], type="value", value='GFPGANv1.4.pth', label='Face Restoration version', info="Face Restoration and RealESR can be freely combined in different ways, or one can be set to \"None\" to use only the other model. Face Restoration is primarily used for face restoration in real-life images, while RealESR serves as a background restoration model.")
641
+ upscale_model = gr.Dropdown(list(typed_upscale_models.keys())+[None], type="value", value='SRVGG, realesr-general-x4v3.pth', label='UpScale version')
642
+ upscale_scale = gr.Number(label="Rescaling factor", value=4)
643
+ face_detection = gr.Dropdown(["retinaface_resnet50", "YOLOv5l", "YOLOv5n"], type="value", value="retinaface_resnet50", label="Face Detection type")
644
+ with_model_name = gr.Checkbox(label="Output image files name with model name", value=True)
645
+ with gr.Row():
646
+ submit = gr.Button(value="Submit", variant="primary", size="lg")
647
+ clear = gr.ClearButton(
648
+ components=[
649
+ input_image,
650
+ face_model,
651
+ upscale_model,
652
+ upscale_scale,
653
+ face_detection,
654
+ with_model_name,
655
+ ], variant="secondary", size="lg",)
656
+ with gr.Column(variant="panel"):
657
+ gallerys = gr.Gallery(type="filepath", label="Output (The whole image)", format="png")
658
+ outputs = gr.File(label="Download the output image")
659
+ with gr.Row(variant="panel"):
660
+ # Generate output array
661
+ output_arr = []
662
+ for file_name in example_list:
663
+ output_arr.append([file_name,])
664
+ gr.Examples(output_arr, inputs=[input_image,], examples_per_page=20)
665
+ with gr.Row(variant="panel"):
666
+ # Convert to Markdown table
667
+ header = "| Face Model Name | Info | Download URL |\n|------------|------|--------------|"
668
+ rows = [
669
+ f"| [{key}]({value[1]}) | " + value[2].replace("\n", "<br>") + f" | [download]({value[0]}) |"
670
+ for key, value in face_models.items()
671
+ ]
672
+ markdown_table = header + "\n" + "\n".join(rows)
673
+ gr.Markdown(value=markdown_table)
674
+ with gr.Row(variant="panel"):
675
+ # Convert to Markdown table
676
+ header = "| Upscale Model Name | Info | Download URL |\n|------------|------|--------------|"
677
+ rows = [
678
+ f"| [{key}]({value[1]}) | " + value[2].replace("\n", "<br>") + f" | [download]({value[0]}) |"
679
+ for key, value in upscale_models.items()
680
+ ]
681
+ markdown_table = header + "\n" + "\n".join(rows)
682
+ gr.Markdown(value=markdown_table)
683
+
684
+ submit.click(
685
+ upscale.inference,
686
+ inputs=[
687
+ input_image,
688
+ face_model,
689
+ upscale_model,
690
+ upscale_scale,
691
+ face_detection,
692
+ with_model_name,
693
+ ],
694
+ outputs=[gallerys, outputs],
695
+ )
696
 
697
+ demo.queue(default_concurrency_limit=1)
698
  demo.launch(inbrowser=True)
699
 
700
 
images/a01.jpg ADDED
images/a02.jpg ADDED
images/a03.jpg ADDED
images/a04.jpg ADDED
images/b01.jpg ADDED
images/b02.jpg ADDED
images/b03.jpg ADDED
images/b04.jpg ADDED
images/b05.jpg ADDED
images/b06.jpg ADDED
images/b07.jpg ADDED
images/b08.jpg ADDED
images/b09.jpg ADDED
images/b10.jpg ADDED
images/b11.jpg ADDED
images/bus.jpg ADDED
images/c01.jpg ADDED
images/c02.jpg ADDED
images/c03.jpg ADDED
images/c04.jpg ADDED
images/c05.jpg ADDED
images/c06.jpg ADDED
images/c07.jpg ADDED
images/c08.jpg ADDED
images/c09.jpg ADDED
images/c10.jpg ADDED
images/zidane.jpg ADDED
requirements.txt CHANGED
@@ -1,6 +1,6 @@
1
  --extra-index-url https://download.pytorch.org/whl/cu124
2
 
3
- gradio==5.8.0
4
 
5
  basicsr @ git+https://github.com/avan06/BasicSR
6
  facexlib @ git+https://github.com/avan06/facexlib
 
1
  --extra-index-url https://download.pytorch.org/whl/cu124
2
 
3
+ gradio==5.9.0
4
 
5
  basicsr @ git+https://github.com/avan06/BasicSR
6
  facexlib @ git+https://github.com/avan06/facexlib