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