From 6a9b33c848281cb02f38764e4f91ef767f5e3edd Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Wed, 7 Sep 2022 12:32:28 +0300
Subject: codeformer support

---
 webui.py | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

(limited to 'webui.py')

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
 
-- 
cgit v1.2.1