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/modelloader.py | 74 ++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 54 insertions(+), 20 deletions(-) (limited to 'modules/modelloader.py') diff --git a/modules/modelloader.py b/modules/modelloader.py index 3bd1de4d..6de65c69 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -1,34 +1,36 @@ import os import shutil +import importlib from urllib.parse import urlparse from basicsr.utils.download_util import load_file_from_url +from modules import shared +from modules.upscaler import Upscaler from modules.paths import script_path, models_path -def load_models(model_path: str, model_url: str = None, command_path: str = None, dl_name: str = None, existing=None, - ext_filter=None) -> list: +def load_models(model_path: str, model_url: str = None, command_path: str = None, ext_filter=None, download_name=None) -> list: """ A one-and done loader to try finding the desired models in specified directories. - @param dl_name: The file name to use for downloading a model. If not specified, it will be used from the URL. - @param model_url: If specified, attempt to download model from the given URL. + @param download_name: Specify to download from model_url immediately. + @param model_url: If no other models are found, this will be downloaded on upscale. @param model_path: The location to store/find models in. @param command_path: A command-line argument to search for models in first. - @param existing: An array of existing model paths. @param ext_filter: An optional list of filename extensions to filter by @return: A list of paths containing the desired model(s) """ + output = [] + if ext_filter is None: ext_filter = [] - if existing is None: - existing = [] try: places = [] if command_path is not None and command_path != model_path: pretrained_path = os.path.join(command_path, 'experiments/pretrained_models') if os.path.exists(pretrained_path): + print(f"Appending path: {pretrained_path}") places.append(pretrained_path) elif os.path.exists(command_path): places.append(command_path) @@ -36,26 +38,24 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None for place in places: if os.path.exists(place): for file in os.listdir(place): - if os.path.isdir(file): + full_path = os.path.join(place, file) + if os.path.isdir(full_path): continue if len(ext_filter) != 0: model_name, extension = os.path.splitext(file) if extension not in ext_filter: continue - if file not in existing: - path = os.path.join(place, file) - existing.append(path) - if model_url is not None and len(existing) == 0: - if dl_name is not None: - model_file = load_file_from_url(url=model_url, model_dir=model_path, file_name=dl_name, progress=True) + if file not in output: + output.append(full_path) + if model_url is not None and len(output) == 0: + if download_name is not None: + dl = load_file_from_url(model_url, model_path, True, download_name) + output.append(dl) else: - model_file = load_file_from_url(url=model_url, model_dir=model_path, progress=True) - - if os.path.exists(model_file) and os.path.isfile(model_file) and model_file not in existing: - existing.append(model_file) + output.append(model_url) except: pass - return existing + return output def friendly_name(file: str): @@ -110,4 +110,38 @@ def move_files(src_path: str, dest_path: str, ext_filter: str = None): print(f"Removing empty folder: {src_path}") shutil.rmtree(src_path, True) except: - pass \ No newline at end of file + pass + + +def load_upscalers(): + datas = [] + for cls in Upscaler.__subclasses__(): + name = cls.__name__ + module_name = cls.__module__ + print(f"Class: {name} and {module_name}") + module = importlib.import_module(module_name) + class_ = getattr(module, name) + cmd_name = f"{name.lower().replace('upscaler', '')}-models-path" + print(f"CMD Name: {cmd_name}") + opt_string = None + try: + opt_string = shared.opts.__getattr__(cmd_name) + except: + pass + scaler = class_(opt_string) + for child in scaler.scalers: + print(f"Appending {child.name}") + datas.append(child) + + shared.sd_upscalers = datas + + # for scaler in subclasses: + # print(f"Found scaler: {type(scaler).__name__}") + # try: + # scaler = scaler() + # for child in scaler.scalers: + # print(f"Appending {child.name}") + # datas.append[child] + # except: + # pass + # shared.sd_upscalers = datas -- cgit v1.2.3