From 3dd5e19c83fd52084776f5fdc7f8cb09bfdb958a Mon Sep 17 00:00:00 2001
From: DenOfEquity <166248528+DenOfEquity@users.noreply.github.com>
Date: Fri, 6 Sep 2024 14:09:09 +0100
Subject: [PATCH] Extras tab: gfpgan and cf fixes (#1722)

* fix GFPGAN to work with visibility < 1
* fix codeformer to work with visibility < 1
* try harder to download GFPGAN model. Old method would download only if there were no .pth models in the GFPGAN directory. If codeformer was used before GFPGAN, the supporting models are already downloaded into the GFPGAN directory.
---
 modules/gfpgan_model.py              | 13 ++++++++++++-
 scripts/postprocessing_codeformer.py |  6 ++++--
 scripts/postprocessing_gfpgan.py     |  6 ++++--
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py
index 01ef899e..dffe1613 100644
--- a/modules/gfpgan_model.py
+++ b/modules/gfpgan_model.py
@@ -41,7 +41,18 @@ class FaceRestorerGFPGAN(face_restoration_utils.CommonFaceRestoration):
                     device=self.get_device(),
                     expected_architecture='GFPGAN',
                 ).model
-        raise ValueError("No GFPGAN model found")
+
+        #   if reach here, model not found. previous code will download it iff there are no models in GFPGAN directory
+        #   this will download it if the supporting models exist
+        try:
+            GFPGANmodel = modelloader.load_file_from_url(model_url, model_dir=self.model_path, file_name=model_download_name)
+            return modelloader.load_spandrel_model(
+                GFPGANmodel,
+                device=self.get_device(),
+                expected_architecture='GFPGAN',
+            ).model
+        except:
+            raise ValueError("No GFPGAN model found")
 
     def restore(self, np_image):
         def restore_face(cropped_face_t):
diff --git a/scripts/postprocessing_codeformer.py b/scripts/postprocessing_codeformer.py
index 53a0cc44..1851d7ae 100644
--- a/scripts/postprocessing_codeformer.py
+++ b/scripts/postprocessing_codeformer.py
@@ -25,11 +25,13 @@ class ScriptPostprocessingCodeFormer(scripts_postprocessing.ScriptPostprocessing
         if codeformer_visibility == 0 or not enable:
             return
 
-        restored_img = codeformer_model.codeformer.restore(np.array(pp.image.convert("RGB"), dtype=np.uint8), w=codeformer_weight)
+        source_img = pp.image.convert("RGB")
+
+        restored_img = codeformer_model.codeformer.restore(np.array(source_img, dtype=np.uint8), w=codeformer_weight)
         res = Image.fromarray(restored_img)
 
         if codeformer_visibility < 1.0:
-            res = Image.blend(pp.image, res, codeformer_visibility)
+            res = Image.blend(source_img, res, codeformer_visibility)
 
         pp.image = res
         pp.info["CodeFormer visibility"] = round(codeformer_visibility, 3)
diff --git a/scripts/postprocessing_gfpgan.py b/scripts/postprocessing_gfpgan.py
index 57e36239..b1a52028 100644
--- a/scripts/postprocessing_gfpgan.py
+++ b/scripts/postprocessing_gfpgan.py
@@ -22,11 +22,13 @@ class ScriptPostprocessingGfpGan(scripts_postprocessing.ScriptPostprocessing):
         if gfpgan_visibility == 0 or not enable:
             return
 
-        restored_img = gfpgan_model.gfpgan_fix_faces(np.array(pp.image.convert("RGB"), dtype=np.uint8))
+        source_img = pp.image.convert("RGB")
+
+        restored_img = gfpgan_model.gfpgan_fix_faces(np.array(source_img, dtype=np.uint8))
         res = Image.fromarray(restored_img)
 
         if gfpgan_visibility < 1.0:
-            res = Image.blend(pp.image, res, gfpgan_visibility)
+            res = Image.blend(source_img, res, gfpgan_visibility)
 
         pp.image = res
         pp.info["GFPGAN visibility"] = round(gfpgan_visibility, 3)