codeformer support

1 files changed, 6 insertions, 6 deletions

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)