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/modelloader.py | 65 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 modules/modelloader.py (limited to 'modules/modelloader.py') diff --git a/modules/modelloader.py b/modules/modelloader.py new file mode 100644 index 00000000..d59fbe05 --- /dev/null +++ b/modules/modelloader.py @@ -0,0 +1,65 @@ +import os +from urllib.parse import urlparse + +from basicsr.utils.download_util import load_file_from_url + + +def load_models(model_path: str, model_url: str = None, command_path: str = None, dl_name: str = None, existing=None, + ext_filter=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 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) + """ + 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): + places.append(pretrained_path) + elif os.path.exists(command_path): + places.append(command_path) + places.append(model_path) + for place in places: + if os.path.exists(place): + for file in os.listdir(place): + if os.path.isdir(file): + 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: + 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) + 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) + except: + pass + return existing + + +def friendly_name(file: str): + if "http" in file: + file = urlparse(file).path + + file = os.path.basename(file) + model_name, extension = os.path.splitext(file) + model_name = model_name.replace("_", " ").title() + return model_name -- cgit v1.2.1 From 7d5c29b674bacc5654f8613af134632b7cbdb158 Mon Sep 17 00:00:00 2001 From: d8ahazard Date: Mon, 26 Sep 2022 10:27:18 -0500 Subject: Cleanup existing directories, fixes --- modules/modelloader.py | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) (limited to 'modules/modelloader.py') diff --git a/modules/modelloader.py b/modules/modelloader.py index d59fbe05..9520a681 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -1,8 +1,11 @@ import os +import shutil from urllib.parse import urlparse from basicsr.utils.download_util import load_file_from_url +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: @@ -63,3 +66,38 @@ def friendly_name(file: str): model_name, extension = os.path.splitext(file) model_name = model_name.replace("_", " ").title() return model_name + + +def cleanup_models(): + root_path = script_path + src_path = os.path.join(root_path, "ESRGAN") + dest_path = os.path.join(models_path, "ESRGAN") + move_files(src_path, dest_path) + src_path = os.path.join(root_path, "gfpgan") + dest_path = os.path.join(models_path, "GFPGAN") + move_files(src_path, dest_path) + src_path = os.path.join(root_path, "SwinIR") + dest_path = os.path.join(models_path, "SwinIR") + move_files(src_path, dest_path) + src_path = os.path.join(root_path, "repositories/latent-diffusion/experiments/pretrained_models/") + dest_path = os.path.join(models_path, "LDSR") + move_files(src_path, dest_path) + + +def move_files(src_path: str, dest_path: str): + try: + if not os.path.exists(dest_path): + os.makedirs(dest_path) + if os.path.exists(src_path): + for file in os.listdir(src_path): + if os.path.isfile(file): + fullpath = os.path.join(src_path, file) + print("Moving file: %s to %s" % (fullpath, dest_path)) + try: + shutil.move(fullpath, dest_path) + except: + pass + print("Removing folder: %s" % src_path) + shutil.rmtree(src_path, True) + except: + pass \ No newline at end of file -- cgit v1.2.1 From 11875f586323cea7c5b8398976449788a83dee76 Mon Sep 17 00:00:00 2001 From: d8ahazard Date: Tue, 27 Sep 2022 11:01:13 -0500 Subject: Use model loader with stable-diffusion too. Hook the model loader into the SD_models file. Add default url/download if checkpoint is not found. Add matching stablediffusion-models-path argument. Add message that --ckpt-dir will be removed in the future, but have it pipe to stablediffusion-models-path for now. Update help strings for models-path args so they're more or less uniform. Move sd_model "setup" call to webUI with the others. Ensure "cleanup_models" method moves existing models to the new locations, including SD, and that we aren't deleting folders that still have stuff in them. --- modules/modelloader.py | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) (limited to 'modules/modelloader.py') diff --git a/modules/modelloader.py b/modules/modelloader.py index 9520a681..2ee364f0 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -45,7 +45,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None if file not in existing: path = os.path.join(place, file) existing.append(path) - if model_url is not None: + 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) else: @@ -69,7 +69,13 @@ def friendly_name(file: str): def cleanup_models(): + # This code could probably be more efficient if we used a tuple list or something to store the src/destinations + # and then enumerate that, but this works for now. In the future, it'd be nice to just have every "model" scaler + # somehow auto-register and just do these things... root_path = script_path + src_path = models_path + dest_path = os.path.join(models_path, "Stable-diffusion") + move_files(src_path, dest_path, ".ckpt") src_path = os.path.join(root_path, "ESRGAN") dest_path = os.path.join(models_path, "ESRGAN") move_files(src_path, dest_path) @@ -84,20 +90,25 @@ def cleanup_models(): move_files(src_path, dest_path) -def move_files(src_path: str, dest_path: str): +def move_files(src_path: str, dest_path: str, ext_filter: str = None): try: if not os.path.exists(dest_path): os.makedirs(dest_path) if os.path.exists(src_path): for file in os.listdir(src_path): - if os.path.isfile(file): - fullpath = os.path.join(src_path, file) - print("Moving file: %s to %s" % (fullpath, dest_path)) + fullpath = os.path.join(src_path, file) + if os.path.isfile(fullpath): + print(f"Checking file {file} in {src_path}") + if ext_filter is not None: + if ext_filter not in file: + continue + print(f"Moving {file} from {src_path} to {dest_path}.") try: shutil.move(fullpath, dest_path) except: pass - print("Removing folder: %s" % src_path) - shutil.rmtree(src_path, True) + if len(os.listdir(src_path)) == 0: + print(f"Removing empty folder: {src_path}") + shutil.rmtree(src_path, True) except: pass \ No newline at end of file -- cgit v1.2.1 From 31ad536c331df14dd785bfd2a1f93f91a8f7839e Mon Sep 17 00:00:00 2001 From: d8ahazard Date: Tue, 27 Sep 2022 11:05:25 -0500 Subject: Remove debugging message --- modules/modelloader.py | 1 - 1 file changed, 1 deletion(-) (limited to 'modules/modelloader.py') diff --git a/modules/modelloader.py b/modules/modelloader.py index 2ee364f0..3bd1de4d 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -98,7 +98,6 @@ def move_files(src_path: str, dest_path: str, ext_filter: str = None): for file in os.listdir(src_path): fullpath = os.path.join(src_path, file) if os.path.isfile(fullpath): - print(f"Checking file {file} in {src_path}") if ext_filter is not None: if ext_filter not in file: continue -- cgit v1.2.1 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.1 From 435fd2112aee9a0e61408ac56663e41beea1e446 Mon Sep 17 00:00:00 2001 From: d8ahazard Date: Thu, 29 Sep 2022 19:59:53 -0500 Subject: Fixes, cleanup. --- modules/modelloader.py | 14 -------------- 1 file changed, 14 deletions(-) (limited to 'modules/modelloader.py') diff --git a/modules/modelloader.py b/modules/modelloader.py index 6de65c69..51b3ecd5 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -118,11 +118,9 @@ def load_upscalers(): 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) @@ -130,18 +128,6 @@ def load_upscalers(): 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.1