codeformer support

1 files changed, 12 insertions, 8 deletions

diff --git a/webui.py b/webui.py
index d20ff38f..c9a800dd 100644
--- a/webui.py
+++ b/webui.py
@@ -19,7 +19,9 @@ from modules.ui import plaintext_to_html
 import modules.scripts

 import modules.processing as processing

 import modules.sd_hijack

-import modules.gfpgan_model as gfpgan

+import modules.codeformer_model

+import modules.gfpgan_model

+import modules.face_restoration

 import modules.realesrgan_model as realesrgan

 import modules.esrgan_model as esrgan

 import modules.images as images

@@ -28,10 +30,12 @@ import modules.txt2img
 import modules.img2img

 

 

+modules.codeformer_model.setup_codeformer()

+modules.gfpgan_model.setup_gfpgan()

+shared.face_restorers.append(modules.face_restoration.FaceRestoration())

+

 esrgan.load_models(cmd_opts.esrgan_models_path)

 realesrgan.setup_realesrgan()

-gfpgan.setup_gfpgan()

-

 

 def load_model_from_config(config, ckpt, verbose=False):

     print(f"Loading model from {ckpt}")

@@ -54,19 +58,19 @@ def load_model_from_config(config, ckpt, verbose=False):
 cached_images = {}

 

 

-def run_extras(image, gfpgan_strength, upscaling_resize, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility):

+def run_extras(image, face_restoration_blending, 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 gfpgan.have_gfpgan is not None and gfpgan_strength > 0:

-        restored_img = gfpgan.gfpgan_fix_faces(np.array(image, dtype=np.uint8))

+    if face_restoration_blending > 0:

+        restored_img = modules.face_restoration.restore_faces(np.array(image, dtype=np.uint8))

         res = Image.fromarray(restored_img)

 

-        if gfpgan_strength < 1.0:

-            res = Image.blend(image, res, gfpgan_strength)

+        if face_restoration_blending < 1.0:

+            res = Image.blend(image, res, face_restoration_blending)

 

         image = res