Harisreedhar commited on
Commit
d8ef00b
β€’
1 Parent(s): db275a2

add codeformer

Browse files
assets/pretrained_models/{codeformer.pth β†’ codeformer.onnx} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:1009e537e0c2a07d4cabce6355f53cb66767cd4b4297ec7a4a64ca4b8a5684b7
3
- size 376637898
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:91e7e881c5001fea4a535e8f96eaeaa672d30c963a678a3e27f0429a6620f57a
3
+ size 376821950
assets/pretrained_models/nsfwmodel_281.pth DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:ac92f5326f0d83f24f51ba4ac9f2a79314d29199e900a8ea495a74816ad3eb67
3
- size 4925
 
 
 
 
face_enhancer.py CHANGED
@@ -4,7 +4,7 @@ import torch
4
  import gfpgan
5
  from PIL import Image
6
  from upscaler.RealESRGAN import RealESRGAN
7
-
8
 
9
  def gfpgan_runner(img, model):
10
  _, imgs, _ = model.enhance(img, paste_back=True, has_aligned=True)
@@ -16,7 +16,13 @@ def realesrgan_runner(img, model):
16
  return img
17
 
18
 
 
 
 
 
 
19
  supported_enhancers = {
 
20
  "GFPGAN": ("./assets/pretrained_models/GFPGANv1.4.pth", gfpgan_runner),
21
  "REAL-ESRGAN 2x": ("./assets/pretrained_models/RealESRGAN_x2.pth", realesrgan_runner),
22
  "REAL-ESRGAN 4x": ("./assets/pretrained_models/RealESRGAN_x4.pth", realesrgan_runner),
@@ -39,7 +45,9 @@ def load_face_enhancer_model(name='GFPGAN', device="cpu"):
39
  if name in supported_enhancers.keys():
40
  model_path, model_runner = supported_enhancers.get(name)
41
  model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
42
- if name == 'GFPGAN':
 
 
43
  model = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=device)
44
  elif name == 'REAL-ESRGAN 2x':
45
  model = RealESRGAN(device, scale=2)
 
4
  import gfpgan
5
  from PIL import Image
6
  from upscaler.RealESRGAN import RealESRGAN
7
+ from upscaler.codeformer import CodeFormerEnhancer
8
 
9
  def gfpgan_runner(img, model):
10
  _, imgs, _ = model.enhance(img, paste_back=True, has_aligned=True)
 
16
  return img
17
 
18
 
19
+ def codeformer_runner(img, model):
20
+ img = model.enhance(img)
21
+ return img
22
+
23
+
24
  supported_enhancers = {
25
+ "CodeFormer": ("./assets/pretrained_models/codeformer.onnx", codeformer_runner),
26
  "GFPGAN": ("./assets/pretrained_models/GFPGANv1.4.pth", gfpgan_runner),
27
  "REAL-ESRGAN 2x": ("./assets/pretrained_models/RealESRGAN_x2.pth", realesrgan_runner),
28
  "REAL-ESRGAN 4x": ("./assets/pretrained_models/RealESRGAN_x4.pth", realesrgan_runner),
 
45
  if name in supported_enhancers.keys():
46
  model_path, model_runner = supported_enhancers.get(name)
47
  model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
48
+ if name == 'CodeFormer':
49
+ model = CodeFormerEnhancer(model_path=model_path, device=device)
50
+ elif name == 'GFPGAN':
51
  model = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=device)
52
  elif name == 'REAL-ESRGAN 2x':
53
  model = RealESRGAN(device, scale=2)
upscaler/codeformer.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import torch
3
+ import onnx
4
+ import onnxruntime
5
+ import numpy as np
6
+
7
+ import time
8
+
9
+ # codeformer converted to onnx
10
+ # using https://github.com/redthing1/CodeFormer
11
+
12
+
13
+ class CodeFormerEnhancer:
14
+ def __init__(self, model_path="CodeFormer/weights/CodeFormer/codeformer.pth", device='cpu'):
15
+ model = onnx.load(model_path)
16
+ session_options = onnxruntime.SessionOptions()
17
+ session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
18
+ providers = ["CPUExecutionProvider"]
19
+ if device == 'cuda':
20
+ providers = [("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}),"CPUExecutionProvider"]
21
+ self.session = onnxruntime.InferenceSession("codeformer.onnx", sess_options=session_options, providers=providers)
22
+
23
+ def enhance(self, img, w=0.9):
24
+ img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
25
+ img = img.astype(np.float32)[:,:,::-1] / 255.0
26
+ img = img.transpose((2, 0, 1))
27
+ nrm_mean = np.array([0.5, 0.5, 0.5]).reshape((-1, 1, 1))
28
+ nrm_std = np.array([0.5, 0.5, 0.5]).reshape((-1, 1, 1))
29
+ img = (img - nrm_mean) / nrm_std
30
+
31
+ img = np.expand_dims(img, axis=0)
32
+
33
+ out = self.session.run(None, {'x':img.astype(np.float32), 'w':np.array([w], dtype=np.double)})[0]
34
+ out = (out[0].transpose(1,2,0).clip(-1,1) + 1) * 0.5
35
+ out = (out * 255)[:,:,::-1]
36
+
37
+ return out.astype('uint8')