victorisgeek commited on
Commit
9bb4c11
1 Parent(s): a67c6b6

Update face_enhancer.py

Browse files
Files changed (1) hide show
  1. face_enhancer.py +6 -6
face_enhancer.py CHANGED
@@ -34,9 +34,9 @@ def codeformer_runner(img, model):
34
  supported_enhancers = {
35
  "CodeFormer": ("./assets/pretrained_models/codeformer.onnx", codeformer_runner),
36
  "GFPGAN": ("./assets/pretrained_models/GFPGANv1.4.pth", gfpgan_runner),
37
- "Real_Hatgan x4": ("./assets/pretrained_models/real_hatgan_x4.onnx", realhatgan_runner),
38
- "Ultra_Sharp x4": ("./assets/pretrained_models/ultra_sharp_x4.onnx", ultrasharp_runner),
39
- "RestoreFormer": ("./assets/pretrained_models/restoreformer_plus_plus.onnx", restoreformer_runner),
40
  "REAL-ESRGAN 2x": ("./assets/pretrained_models/RealESRGAN_x2.pth", realesrgan_runner),
41
  "REAL-ESRGAN 4x": ("./assets/pretrained_models/RealESRGAN_x4.pth", realesrgan_runner),
42
  "REAL-ESRGAN 8x": ("./assets/pretrained_models/RealESRGAN_x8.pth", realesrgan_runner)
@@ -62,7 +62,7 @@ def load_face_enhancer_model(name='GFPGAN', device="cpu"):
62
  model = CodeFormerEnhancer(model_path=model_path, device=device)
63
  elif name == 'GFPGAN':
64
  model = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=device)
65
- elif name =='RestoreFormer':
66
  model = RestoreFormer(device, scale=1)
67
  model.load_weights(model_path, download=False)
68
  elif name == 'REAL-ESRGAN 4x':
@@ -71,9 +71,9 @@ def load_face_enhancer_model(name='GFPGAN', device="cpu"):
71
  elif name == 'REAL-ESRGAN 2x':
72
  model = RealESRGAN(device, scale=2)
73
  model.load_weights(model_path, download=False)
74
- elif name == 'Real_Hatgan x4':
75
  model = RealHatGAN(device, scale=4)
76
- elif name == 'Ultra_Sharp x4':
77
  model = UltraShap(device, scale=4)
78
  model.load_weights(model_path, download=False)
79
  elif name == 'REAL-ESRGAN 4x':
 
34
  supported_enhancers = {
35
  "CodeFormer": ("./assets/pretrained_models/codeformer.onnx", codeformer_runner),
36
  "GFPGAN": ("./assets/pretrained_models/GFPGANv1.4.pth", gfpgan_runner),
37
+ "REAL_HATGAN x4": ("./assets/pretrained_models/real_hatgan_x4.onnx", realhatgan_runner),
38
+ "ULTRA_SHARP x4": ("./assets/pretrained_models/ultra_sharp_x4.onnx", ultrasharp_runner),
39
+ "RESTOREFORMER PLUS PLUS": ("./assets/pretrained_models/restoreformer_plus_plus.onnx", restoreformer_runner),
40
  "REAL-ESRGAN 2x": ("./assets/pretrained_models/RealESRGAN_x2.pth", realesrgan_runner),
41
  "REAL-ESRGAN 4x": ("./assets/pretrained_models/RealESRGAN_x4.pth", realesrgan_runner),
42
  "REAL-ESRGAN 8x": ("./assets/pretrained_models/RealESRGAN_x8.pth", realesrgan_runner)
 
62
  model = CodeFormerEnhancer(model_path=model_path, device=device)
63
  elif name == 'GFPGAN':
64
  model = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=device)
65
+ elif name =='RESTOREFORMER PLUS PLUS':
66
  model = RestoreFormer(device, scale=1)
67
  model.load_weights(model_path, download=False)
68
  elif name == 'REAL-ESRGAN 4x':
 
71
  elif name == 'REAL-ESRGAN 2x':
72
  model = RealESRGAN(device, scale=2)
73
  model.load_weights(model_path, download=False)
74
+ elif name == 'REAL_HATGAN x4':
75
  model = RealHatGAN(device, scale=4)
76
+ elif name == 'ULTRA_SHARP x4':
77
  model = UltraShap(device, scale=4)
78
  model.load_weights(model_path, download=False)
79
  elif name == 'REAL-ESRGAN 4x':