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

---
 modules/processing.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

(limited to 'modules/processing.py')

diff --git a/modules/processing.py b/modules/processing.py
index 78bc73b9..49474b73 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -14,7 +14,7 @@ from modules.sd_hijack import model_hijack
 from modules.sd_samplers import samplers, samplers_for_img2img
 from modules.shared import opts, cmd_opts, state
 import modules.shared as shared
-import modules.gfpgan_model as gfpgan
+import modules.face_restoration
 import modules.images as images
 
 # some of those options should not be changed at all because they would break the model, so I removed them from options.
@@ -29,7 +29,7 @@ def torch_gc():
 
 
 class StableDiffusionProcessing:
-    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt="", seed=-1, sampler_index=0, batch_size=1, n_iter=1, steps=50, cfg_scale=7.0, width=512, height=512, use_GFPGAN=False, tiling=False, do_not_save_samples=False, do_not_save_grid=False, extra_generation_params=None, overlay_images=None, negative_prompt=None):
+    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt="", seed=-1, sampler_index=0, batch_size=1, n_iter=1, steps=50, cfg_scale=7.0, width=512, height=512, restore_faces=False, tiling=False, do_not_save_samples=False, do_not_save_grid=False, extra_generation_params=None, overlay_images=None, negative_prompt=None):
         self.sd_model = sd_model
         self.outpath_samples: str = outpath_samples
         self.outpath_grids: str = outpath_grids
@@ -44,7 +44,7 @@ class StableDiffusionProcessing:
         self.cfg_scale: float = cfg_scale
         self.width: int = width
         self.height: int = height
-        self.use_GFPGAN: bool = use_GFPGAN
+        self.restore_faces: bool = restore_faces
         self.tiling: bool = tiling
         self.do_not_save_samples: bool = do_not_save_samples
         self.do_not_save_grid: bool = do_not_save_grid
@@ -136,7 +136,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             "Sampler": samplers[p.sampler_index].name,
             "CFG scale": p.cfg_scale,
             "Seed": all_seeds[position_in_batch + iteration * p.batch_size],
-            "GFPGAN": ("GFPGAN" if p.use_GFPGAN else None),
+            "Face restoration": (opts.face_restoration_model if p.restore_faces else None),
             "Batch size": (None if p.batch_size < 2 else p.batch_size),
             "Batch pos": (None if p.batch_size < 2 else position_in_batch),
         }
@@ -193,10 +193,10 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
                 x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
                 x_sample = x_sample.astype(np.uint8)
 
-                if p.use_GFPGAN:
+                if p.restore_faces:
                     torch_gc()
 
-                    x_sample = gfpgan.gfpgan_fix_faces(x_sample)
+                    x_sample = modules.face_restoration.restore_faces(x_sample)
 
                 image = Image.fromarray(x_sample)
 
-- 
cgit v1.2.1