diff options
| author | AUTOMATIC <16777216c@gmail.com> | 2022-09-03 09:08:45 +0000 |
|---|---|---|
| committer | AUTOMATIC <16777216c@gmail.com> | 2022-09-03 09:08:45 +0000 |
| commit | 345028099d893f8a66726cfd13627d8cc1bcc724 (patch) | |
| tree | acb1da553620b7e7139db840ef43accf71b786a8 /modules/gfpgan_model.py | |
| parent | d7b67d9b40e47ede766d3beb149b0c2b74651ece (diff) | |
| download | stable-diffusion-webui-gfx803-345028099d893f8a66726cfd13627d8cc1bcc724.tar.gz stable-diffusion-webui-gfx803-345028099d893f8a66726cfd13627d8cc1bcc724.tar.bz2 stable-diffusion-webui-gfx803-345028099d893f8a66726cfd13627d8cc1bcc724.zip | |
split codebase into multiple files; to anyone this affects negatively: sorry
Diffstat (limited to 'modules/gfpgan_model.py')
| -rw-r--r-- | modules/gfpgan_model.py | 58 |
1 files changed, 58 insertions, 0 deletions
diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py new file mode 100644 index 00000000..3f42c163 --- /dev/null +++ b/modules/gfpgan_model.py @@ -0,0 +1,58 @@ +import os
+import sys
+import traceback
+
+from modules.paths import script_path
+from modules.shared import cmd_opts
+
+
+def gfpgan_model_path():
+ places = [script_path, '.', os.path.join(cmd_opts.gfpgan_dir, 'experiments/pretrained_models')]
+ files = [cmd_opts.gfpgan_model] + [os.path.join(dirname, cmd_opts.gfpgan_model) for dirname in places]
+ found = [x for x in files if os.path.exists(x)]
+
+ if len(found) == 0:
+ raise Exception("GFPGAN model not found in paths: " + ", ".join(files))
+
+ return found[0]
+
+
+loaded_gfpgan_model = None
+
+
+def gfpgan():
+ global loaded_gfpgan_model
+
+ if loaded_gfpgan_model is None and gfpgan_constructor is not None:
+ loaded_gfpgan_model = gfpgan_constructor(model_path=gfpgan_model_path(), upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None)
+
+ return loaded_gfpgan_model
+
+
+def gfpgan_fix_faces(np_image):
+ np_image_bgr = np_image[:, :, ::-1]
+ cropped_faces, restored_faces, gfpgan_output_bgr = gfpgan().enhance(np_image_bgr, has_aligned=False, only_center_face=False, paste_back=True)
+ np_image = gfpgan_output_bgr[:, :, ::-1]
+
+ return np_image
+
+
+have_gfpgan = False
+gfpgan_constructor = None
+
+def setup_gfpgan():
+ try:
+ gfpgan_model_path()
+
+ if os.path.exists(cmd_opts.gfpgan_dir):
+ sys.path.append(os.path.abspath(cmd_opts.gfpgan_dir))
+ from gfpgan import GFPGANer
+
+ global have_gfpgan
+ have_gfpgan = True
+
+ global gfpgan_constructor
+ gfpgan_constructor = GFPGANer
+ except Exception:
+ print("Error setting up GFPGAN:", file=sys.stderr)
+ print(traceback.format_exc(), file=sys.stderr)
|