From 740070ea9cdb254209f66417418f2a4af8b099d6 Mon Sep 17 00:00:00 2001 From: d8ahazard Date: Mon, 26 Sep 2022 09:29:50 -0500 Subject: Re-implement universal model loading --- modules/gfpgan_model.py | 60 ++++++++++++++++++++++++------------------------- 1 file changed, 29 insertions(+), 31 deletions(-) (limited to 'modules/gfpgan_model.py') diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py index 44c5dc6c..ffb6960d 100644 --- a/modules/gfpgan_model.py +++ b/modules/gfpgan_model.py @@ -7,33 +7,20 @@ from modules import shared, devices from modules.shared import cmd_opts from modules.paths import script_path import modules.face_restoration +from modules import shared, devices, modelloader +from modules.paths import models_path - -def gfpgan_model_path(): - from modules.shared import cmd_opts - - filemask = 'GFPGAN*.pth' - - if cmd_opts.gfpgan_model is not None: - return cmd_opts.gfpgan_model - - places = [script_path, '.', os.path.join(cmd_opts.gfpgan_dir, 'experiments/pretrained_models')] - - filename = None - for place in places: - filename = next(iter(glob(os.path.join(place, filemask))), None) - if filename is not None: - break - - return filename - +model_dir = "GFPGAN" +cmd_dir = None +model_path = os.path.join(models_path, model_dir) +model_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth" loaded_gfpgan_model = None def gfpgan(): global loaded_gfpgan_model - + global model_path if loaded_gfpgan_model is not None: loaded_gfpgan_model.gfpgan.to(shared.device) return loaded_gfpgan_model @@ -41,7 +28,15 @@ def gfpgan(): if gfpgan_constructor is None: return None - model = gfpgan_constructor(model_path=gfpgan_model_path() or 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth', upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None) + models = modelloader.load_models(model_path, model_url, cmd_dir) + if len(models) != 0: + latest_file = max(models, key=os.path.getctime) + model_file = latest_file + else: + print("Unable to load gfpgan model!") + return None + model = gfpgan_constructor(model_path=model_file, model_dir=model_path, upscale=1, arch='clean', channel_multiplier=2, + bg_upsampler=None) model.gfpgan.to(shared.device) loaded_gfpgan_model = model @@ -50,7 +45,8 @@ def gfpgan(): def gfpgan_fix_faces(np_image): model = gfpgan() - + if model is None: + return np_image np_image_bgr = np_image[:, :, ::-1] cropped_faces, restored_faces, gfpgan_output_bgr = model.enhance(np_image_bgr, has_aligned=False, only_center_face=False, paste_back=True) np_image = gfpgan_output_bgr[:, :, ::-1] @@ -64,19 +60,21 @@ def gfpgan_fix_faces(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 +def setup_model(dirname): + global model_path + if not os.path.exists(model_path): + os.makedirs(model_path) + try: + from modules.gfpgan_model_arch import GFPGANerr + global cmd_dir global have_gfpgan - have_gfpgan = True - global gfpgan_constructor - gfpgan_constructor = GFPGANer + + cmd_dir = dirname + have_gfpgan = True + gfpgan_constructor = GFPGANerr class FaceRestorerGFPGAN(modules.face_restoration.FaceRestoration): def name(self): -- cgit v1.2.3 From 0dce0df1ee63b2f158805c1a1f1a3743cc4a104b Mon Sep 17 00:00:00 2001 From: d8ahazard Date: Thu, 29 Sep 2022 17:46:23 -0500 Subject: Holy $hit. Yep. Fix gfpgan_model_arch requirement(s). Add Upscaler base class, move from images. Add a lot of methods to Upscaler. Re-work all the child upscalers to be proper classes. Add BSRGAN scaler. Add ldsr_model_arch class, removing the dependency for another repo that just uses regular latent-diffusion stuff. Add one universal method that will always find and load new upscaler models without having to add new "setup_model" calls. Still need to add command line params, but that could probably be automated. Add a "self.scale" property to all Upscalers so the scalers themselves can do "things" in response to the requested upscaling size. Ensure LDSR doesn't get stuck in a longer loop of "upscale/downscale/upscale" as we try to reach the target upscale size. Add typehints for IDE sanity. PEP-8 improvements. Moar. --- modules/gfpgan_model.py | 58 ++++++++++++++++++++++++++++++++++--------------- 1 file changed, 40 insertions(+), 18 deletions(-) (limited to 'modules/gfpgan_model.py') diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py index ffb6960d..2bf8a1ee 100644 --- a/modules/gfpgan_model.py +++ b/modules/gfpgan_model.py @@ -1,24 +1,23 @@ import os import sys import traceback -from glob import glob -from modules import shared, devices -from modules.shared import cmd_opts -from modules.paths import script_path +import facexlib +import gfpgan + import modules.face_restoration from modules import shared, devices, modelloader from modules.paths import models_path model_dir = "GFPGAN" -cmd_dir = None +user_path = None model_path = os.path.join(models_path, model_dir) model_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth" - +have_gfpgan = False loaded_gfpgan_model = None -def gfpgan(): +def gfpgann(): global loaded_gfpgan_model global model_path if loaded_gfpgan_model is not None: @@ -28,14 +27,16 @@ def gfpgan(): if gfpgan_constructor is None: return None - models = modelloader.load_models(model_path, model_url, cmd_dir) - if len(models) != 0: + models = modelloader.load_models(model_path, model_url, user_path, ext_filter="GFPGAN") + if len(models) == 1 and "http" in models[0]: + model_file = models[0] + elif len(models) != 0: latest_file = max(models, key=os.path.getctime) model_file = latest_file else: print("Unable to load gfpgan model!") return None - model = gfpgan_constructor(model_path=model_file, model_dir=model_path, upscale=1, arch='clean', channel_multiplier=2, + model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None) model.gfpgan.to(shared.device) loaded_gfpgan_model = model @@ -44,11 +45,12 @@ def gfpgan(): def gfpgan_fix_faces(np_image): - model = gfpgan() + model = gfpgann() if model is None: return np_image np_image_bgr = np_image[:, :, ::-1] - cropped_faces, restored_faces, gfpgan_output_bgr = model.enhance(np_image_bgr, has_aligned=False, only_center_face=False, paste_back=True) + cropped_faces, restored_faces, gfpgan_output_bgr = model.enhance(np_image_bgr, has_aligned=False, + only_center_face=False, paste_back=True) np_image = gfpgan_output_bgr[:, :, ::-1] if shared.opts.face_restoration_unload: @@ -57,7 +59,6 @@ def gfpgan_fix_faces(np_image): return np_image -have_gfpgan = False gfpgan_constructor = None @@ -67,14 +68,33 @@ def setup_model(dirname): os.makedirs(model_path) try: - from modules.gfpgan_model_arch import GFPGANerr - global cmd_dir + from gfpgan import GFPGANer + from facexlib import detection, parsing + global user_path global have_gfpgan global gfpgan_constructor - cmd_dir = dirname + load_file_from_url_orig = gfpgan.utils.load_file_from_url + facex_load_file_from_url_orig = facexlib.detection.load_file_from_url + facex_load_file_from_url_orig2 = facexlib.parsing.load_file_from_url + + def my_load_file_from_url(**kwargs): + print("Setting model_dir to " + model_path) + return load_file_from_url_orig(**dict(kwargs, model_dir=model_path)) + + def facex_load_file_from_url(**kwargs): + return facex_load_file_from_url_orig(**dict(kwargs, save_dir=model_path, model_dir=None)) + + def facex_load_file_from_url2(**kwargs): + return facex_load_file_from_url_orig2(**dict(kwargs, save_dir=model_path, model_dir=None)) + + gfpgan.utils.load_file_from_url = my_load_file_from_url + facexlib.detection.load_file_from_url = facex_load_file_from_url + facexlib.parsing.load_file_from_url = facex_load_file_from_url2 + user_path = dirname + print("Have gfpgan should be true?") have_gfpgan = True - gfpgan_constructor = GFPGANerr + gfpgan_constructor = GFPGANer class FaceRestorerGFPGAN(modules.face_restoration.FaceRestoration): def name(self): @@ -82,7 +102,9 @@ def setup_model(dirname): def restore(self, 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) + cropped_faces, restored_faces, gfpgan_output_bgr = gfpgann().enhance(np_image_bgr, has_aligned=False, + only_center_face=False, + paste_back=True) np_image = gfpgan_output_bgr[:, :, ::-1] return np_image -- cgit v1.2.3 From d1f098540ad1dbc2abb8d04322634efba650b631 Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Fri, 30 Sep 2022 11:42:40 +0300 Subject: remove unwanted formatting/functionality from the PR --- modules/gfpgan_model.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) (limited to 'modules/gfpgan_model.py') diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py index 2bf8a1ee..bb30d733 100644 --- a/modules/gfpgan_model.py +++ b/modules/gfpgan_model.py @@ -36,8 +36,7 @@ def gfpgann(): else: print("Unable to load gfpgan model!") return None - model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, - bg_upsampler=None) + model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None) model.gfpgan.to(shared.device) loaded_gfpgan_model = model @@ -49,8 +48,7 @@ def gfpgan_fix_faces(np_image): if model is None: return np_image np_image_bgr = np_image[:, :, ::-1] - cropped_faces, restored_faces, gfpgan_output_bgr = model.enhance(np_image_bgr, has_aligned=False, - only_center_face=False, paste_back=True) + cropped_faces, restored_faces, gfpgan_output_bgr = model.enhance(np_image_bgr, has_aligned=False, only_center_face=False, paste_back=True) np_image = gfpgan_output_bgr[:, :, ::-1] if shared.opts.face_restoration_unload: @@ -79,7 +77,6 @@ def setup_model(dirname): facex_load_file_from_url_orig2 = facexlib.parsing.load_file_from_url def my_load_file_from_url(**kwargs): - print("Setting model_dir to " + model_path) return load_file_from_url_orig(**dict(kwargs, model_dir=model_path)) def facex_load_file_from_url(**kwargs): @@ -92,7 +89,6 @@ def setup_model(dirname): facexlib.detection.load_file_from_url = facex_load_file_from_url facexlib.parsing.load_file_from_url = facex_load_file_from_url2 user_path = dirname - print("Have gfpgan should be true?") have_gfpgan = True gfpgan_constructor = GFPGANer @@ -102,9 +98,7 @@ def setup_model(dirname): def restore(self, np_image): np_image_bgr = np_image[:, :, ::-1] - cropped_faces, restored_faces, gfpgan_output_bgr = gfpgann().enhance(np_image_bgr, has_aligned=False, - only_center_face=False, - paste_back=True) + cropped_faces, restored_faces, gfpgan_output_bgr = gfpgann().enhance(np_image_bgr, has_aligned=False, only_center_face=False, paste_back=True) np_image = gfpgan_output_bgr[:, :, ::-1] return np_image -- cgit v1.2.3