From 8dcd3e57e7f56ea45ac815ce479ddfc70122ee27 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Wed, 7 Sep 2022 13:35:02 +0300
Subject: [PATCH] a little bit of rework for extras tab

---
 modules/codeformer_model.py | 10 +++++++---
 modules/ui.py               | 13 ++++++++++---
 webui.py                    | 19 ++++++++++++++-----
 3 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/modules/codeformer_model.py b/modules/codeformer_model.py
index 836417c2..82b45b15 100644
--- a/modules/codeformer_model.py
+++ b/modules/codeformer_model.py
@@ -16,6 +16,7 @@ from importlib import reload
 pretrain_model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
 
 have_codeformer = False
+codeformer = None
 
 def setup_codeformer():
     path = modules.paths.paths.get("CodeFormer", None)
@@ -64,7 +65,7 @@ def setup_codeformer():
 
                 return net, face_helper
 
-            def restore(self, np_image):
+            def restore(self, np_image, w=None):
                 np_image = np_image[:, :, ::-1]
 
                 net, face_helper = self.create_models()
@@ -80,7 +81,7 @@ def setup_codeformer():
 
                     try:
                         with torch.no_grad():
-                            output = net(cropped_face_t, w=shared.opts.code_former_weight, adain=True)[0]
+                            output = net(cropped_face_t, w=w if w is not None else shared.opts.code_former_weight, adain=True)[0]
                             restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
                         del output
                         torch.cuda.empty_cache()
@@ -99,7 +100,10 @@ def setup_codeformer():
 
         global have_codeformer
         have_codeformer = True
-        shared.face_restorers.append(FaceRestorerCodeFormer())
+
+        global codeformer
+        codeformer = FaceRestorerCodeFormer()
+        shared.face_restorers.append(codeformer)
 
     except Exception:
         print("Error setting up CodeFormer:", file=sys.stderr)
diff --git a/modules/ui.py b/modules/ui.py
index 05161b00..689debc9 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -21,9 +21,10 @@ from modules.paths import script_path
 from modules.shared import opts, cmd_opts
 import modules.shared as shared
 from modules.sd_samplers import samplers, samplers_for_img2img
-import modules.gfpgan_model as gfpgan
 import modules.realesrgan_model as realesrgan
 import modules.scripts
+import modules.gfpgan_model
+import modules.codeformer_model
 
 # this is a fix for Windows users. Without it, javascript files will be served with text/html content-type and the bowser will not show any UI
 mimetypes.init()
@@ -521,7 +522,11 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                     extras_upscaler_2_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Upscaler 2 visibility", value=1)
 
                 with gr.Group():
-                    face_restoration_blending = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Faces restoration visibility", value=0, interactive=len(shared.face_restorers) > 1)
+                    gfpgan_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="GFPGAN visibility", value=0, interactive=modules.gfpgan_model.have_gfpgan)
+
+                with gr.Group():
+                    codeformer_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="CodeFormer visibility", value=0, interactive=modules.codeformer_model.have_codeformer)
+                    codeformer_weight = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="CodeFormer weight (0=max, 1=none)", value=0, interactive=modules.codeformer_model.have_codeformer)
 
                 submit = gr.Button('Generate', elem_id="extras_generate", variant='primary')
 
@@ -534,7 +539,9 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
             fn=run_extras,
             inputs=[
                 image,
-                face_restoration_blending,
+                gfpgan_visibility,
+                codeformer_visibility,
+                codeformer_weight,
                 upscaling_resize,
                 extras_upscaler_1,
                 extras_upscaler_2,
diff --git a/webui.py b/webui.py
index c9a800dd..4f701fad 100644
--- a/webui.py
+++ b/webui.py
@@ -58,19 +58,28 @@ def load_model_from_config(config, ckpt, verbose=False):
 cached_images = {}
 
 
-def run_extras(image, face_restoration_blending, upscaling_resize, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility):
+def run_extras(image, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility):
     processing.torch_gc()
 
     image = image.convert("RGB")
 
     outpath = opts.outdir_samples or opts.outdir_extras_samples
 
-    if face_restoration_blending > 0:
-        restored_img = modules.face_restoration.restore_faces(np.array(image, dtype=np.uint8))
+    if gfpgan_visibility > 0:
+        restored_img = modules.gfpgan_model.gfpgan_fix_faces(np.array(image, dtype=np.uint8))
         res = Image.fromarray(restored_img)
 
-        if face_restoration_blending < 1.0:
-            res = Image.blend(image, res, face_restoration_blending)
+        if gfpgan_visibility < 1.0:
+            res = Image.blend(image, res, gfpgan_visibility)
+
+        image = res
+
+    if codeformer_visibility > 0:
+        restored_img = modules.codeformer_model.codeformer.restore(np.array(image, dtype=np.uint8), w=codeformer_weight)
+        res = Image.fromarray(restored_img)
+
+        if codeformer_visibility < 1.0:
+            res = Image.blend(image, res, codeformer_visibility)
 
         image = res