a little bit of rework for extras tab

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