add codeformer

assets/pretrained_models/{codeformer.pth → codeformer.onnx}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1009e537e0c2a07d4cabce6355f53cb66767cd4b4297ec7a4a64ca4b8a5684b7
-size 376637898
+oid sha256:91e7e881c5001fea4a535e8f96eaeaa672d30c963a678a3e27f0429a6620f57a
+size 376821950

assets/pretrained_models/nsfwmodel_281.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ac92f5326f0d83f24f51ba4ac9f2a79314d29199e900a8ea495a74816ad3eb67
-size 4925

face_enhancer.py

@@ -4,7 +4,7 @@ import torch
 import gfpgan
 from PIL import Image
 from upscaler.RealESRGAN import RealESRGAN
-
+from upscaler.codeformer import CodeFormerEnhancer
 
 def gfpgan_runner(img, model):
     _, imgs, _ = model.enhance(img, paste_back=True, has_aligned=True)
@@ -16,7 +16,13 @@ def realesrgan_runner(img, model):
     return img
 
 
+def codeformer_runner(img, model):
+    img = model.enhance(img)
+    return img
+
+
 supported_enhancers = {
+    "CodeFormer": ("./assets/pretrained_models/codeformer.onnx", codeformer_runner),
     "GFPGAN": ("./assets/pretrained_models/GFPGANv1.4.pth", gfpgan_runner),
     "REAL-ESRGAN 2x": ("./assets/pretrained_models/RealESRGAN_x2.pth", realesrgan_runner),
     "REAL-ESRGAN 4x": ("./assets/pretrained_models/RealESRGAN_x4.pth", realesrgan_runner),
@@ -39,7 +45,9 @@ def load_face_enhancer_model(name='GFPGAN', device="cpu"):
     if name in supported_enhancers.keys():
         model_path, model_runner = supported_enhancers.get(name)
         model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model_path)
-    if name == 'GFPGAN':
+    if name == 'CodeFormer':
+        model = CodeFormerEnhancer(model_path=model_path, device=device)
+    elif name == 'GFPGAN':
         model = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=device)
     elif name == 'REAL-ESRGAN 2x':
         model = RealESRGAN(device, scale=2)

upscaler/codeformer.py

@@ -0,0 +1,37 @@
+import cv2
+import torch
+import onnx
+import onnxruntime
+import numpy as np
+
+import time
+
+# codeformer converted to onnx
+# using https://github.com/redthing1/CodeFormer
+
+
+class CodeFormerEnhancer:
+    def __init__(self, model_path="CodeFormer/weights/CodeFormer/codeformer.pth", device='cpu'):
+        model = onnx.load(model_path)
+        session_options = onnxruntime.SessionOptions()
+        session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+        providers = ["CPUExecutionProvider"]
+        if device == 'cuda':
+            providers = [("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}),"CPUExecutionProvider"]
+        self.session = onnxruntime.InferenceSession("codeformer.onnx", sess_options=session_options, providers=providers)
+
+    def enhance(self, img, w=0.9):
+        img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
+        img = img.astype(np.float32)[:,:,::-1] / 255.0
+        img = img.transpose((2, 0, 1))
+        nrm_mean = np.array([0.5, 0.5, 0.5]).reshape((-1, 1, 1))
+        nrm_std = np.array([0.5, 0.5, 0.5]).reshape((-1, 1, 1))
+        img = (img - nrm_mean) / nrm_std
+
+        img = np.expand_dims(img, axis=0)
+
+        out = self.session.run(None, {'x':img.astype(np.float32), 'w':np.array([w], dtype=np.double)})[0]
+        out = (out[0].transpose(1,2,0).clip(-1,1) + 1) * 0.5
+        out = (out * 255)[:,:,::-1]
+
+        return out.astype('uint8')