diff --git a/folder_paths.py b/folder_paths.py index af56a6da1..751b28d5c 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -1,4 +1,5 @@ import os +from collections import namedtuple supported_ckpt_extensions = set(['.ckpt', '.pth']) supported_pt_extensions = set(['.ckpt', '.pt', '.bin', '.pth']) @@ -10,31 +11,24 @@ except: print("Could not import safetensors, safetensors support disabled.") -folder_names_and_paths = {} +SearchLocation = namedtuple('SearchLocation', ['paths', 'extensions']) +model_search_locations = {} -models_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "models") -folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_ckpt_extensions) -folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"]) - -folder_names_and_paths["loras"] = ([os.path.join(models_dir, "loras")], supported_pt_extensions) -folder_names_and_paths["vae"] = ([os.path.join(models_dir, "vae")], supported_pt_extensions) -folder_names_and_paths["clip"] = ([os.path.join(models_dir, "clip")], supported_pt_extensions) -folder_names_and_paths["clip_vision"] = ([os.path.join(models_dir, "clip_vision")], supported_pt_extensions) -folder_names_and_paths["style_models"] = ([os.path.join(models_dir, "style_models")], supported_pt_extensions) -folder_names_and_paths["embeddings"] = ([os.path.join(models_dir, "embeddings")], supported_pt_extensions) - -folder_names_and_paths["controlnet"] = ([os.path.join(models_dir, "controlnet"), os.path.join(models_dir, "t2i_adapter")], supported_pt_extensions) -folder_names_and_paths["upscale_models"] = ([os.path.join(models_dir, "upscale_models")], supported_pt_extensions) +def register_model_type(model_tag, supported_extensions): + if model_tag not in model_search_locations: + model_search_locations[model_tag] = SearchLocation([], supported_extensions) -def add_model_folder_path(folder_name, full_folder_path): - global folder_names_and_paths - if folder_name in folder_names_and_paths: - folder_names_and_paths[folder_name][0].append(full_folder_path) +def add_model_search_path(model_tag, full_folder_path): + global model_search_locations + if model_tag in model_search_locations: + model_search_locations[model_tag].paths.append(full_folder_path) + + +def get_folder_paths(model_tag): + return model_search_locations[model_tag][0][:] -def get_folder_paths(folder_name): - return folder_names_and_paths[folder_name][0][:] def recursive_search(directory): result = [] @@ -44,26 +38,50 @@ def recursive_search(directory): result.append(os.path.join(root, filepath).replace(os.path.join(directory,''),'')) return result + def filter_files_extensions(files, extensions): return sorted(list(filter(lambda a: os.path.splitext(a)[-1].lower() in extensions, files))) - -def get_full_path(folder_name, filename): - global folder_names_and_paths - folders = folder_names_and_paths[folder_name] +def get_full_path(model_tag, filename): + global model_search_locations + folders = model_search_locations[model_tag] for x in folders[0]: full_path = os.path.join(x, filename) if os.path.isfile(full_path): return full_path -def get_filename_list(folder_name): - global folder_names_and_paths +def get_filename_list(model_tag): + global model_search_locations output_list = set() - folders = folder_names_and_paths[folder_name] + folders = model_search_locations[model_tag] for x in folders[0]: output_list.update(filter_files_extensions(recursive_search(x), folders[1])) return sorted(list(output_list)) +# default model types +register_model_type('ckpt', supported_ckpt_extensions) +register_model_type('config', ['.yaml']) +register_model_type('lora', supported_pt_extensions) +register_model_type('vae', supported_pt_extensions) +register_model_type('clip', supported_pt_extensions) +register_model_type('clip_vision', supported_pt_extensions) +register_model_type('style_model', supported_pt_extensions) +register_model_type('embedding', supported_pt_extensions) +register_model_type('controlnet', supported_pt_extensions) +register_model_type('upscale_model', supported_pt_extensions) + +# default search paths +models_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "models") +add_model_search_path('ckpt', os.path.join(models_dir, "checkpoints")) +add_model_search_path('config', os.path.join(models_dir, "configs")) +add_model_search_path('lora', os.path.join(models_dir, "loras")) +add_model_search_path('vae', os.path.join(models_dir, "vae")) +add_model_search_path('clip', os.path.join(models_dir, "clip")) +add_model_search_path('clip_vision', os.path.join(models_dir, "clip_vision")) +add_model_search_path('style_model', os.path.join(models_dir, "style_models")) +add_model_search_path('embedding', os.path.join(models_dir, "embeddings")) +add_model_search_path('controlnet', os.path.join(models_dir, "controlnet")) +add_model_search_path('upscale_model', os.path.join(models_dir, "upscale_models")) diff --git a/main.py b/main.py index e86de8d88..464c7016c 100644 --- a/main.py +++ b/main.py @@ -15,17 +15,17 @@ if __name__ == "__main__": print("\t--listen\t\t\tListen on 0.0.0.0 so the UI can be accessed from other computers.") print("\t--port 8188\t\t\tSet the listen port.") - s = os.path.sep - print(f"\t--ckpt-dir path{s}to{s}dir\t\t\tAdd a path to a checkpoints dir.") - print(f"\t--clip-dir path{s}to{s}dir\t\t\tAdd a path to a clip dir.") - print(f"\t--clip-vision-dir path{s}to{s}dir\tAdd a path to a clip vision dir.") - print(f"\t--controlnet-dir path{s}to{s}dir\tAdd a path to a controlnet checkpoints dir.") - print(f"\t--embed-dir path{s}to{s}dir\t\t\tAdd a path to an embeddings dir.") - print(f"\t--lora-dir path{s}to{s}dir\t\t\tAdd a path to a lora models dir.") - print(f"\t--style-model-dir path{s}to{s}dir\tAdd a path to a style models dir.") - print(f"\t--t2i-dir path{s}to{s}dir\t\t\tAdd a path to a T2I style adapter dir.") - print(f"\t--upscaler-dir path{s}to{s}dir\t\tAdd a path to an upscale models dir.") - print(f"\t--vae-dir path{s}to{s}dir\t\t\tAdd a path to a vae dir.") + dummy_path = os.path.join('path', 'to', 'dir') + print(f"\t--ckpt-dir {dummy_path}\t\t\tAdd a path to a checkpoint directory.") + print(f"\t--config-dir {dummy_path}\t\t\tAdd a path to a model config directory.") + print(f"\t--clip-dir {dummy_path}\t\t\tAdd a path to a clip directory.") + print(f"\t--clip-vision-dir {dummy_path}\tAdd a path to a clip vision directory.") + print(f"\t--controlnet-dir {dummy_path}\tAdd a path to a controlnet or T2I checkpoint directory.") + print(f"\t--embedding-dir {dummy_path}\t\t\tAdd a path to an embedding directory.") + print(f"\t--lora-dir {dummy_path}\t\t\tAdd a path to a lora model directory.") + print(f"\t--style-model-dir {dummy_path}\tAdd a path to a style model directory.") + print(f"\t--upscale-model-dir {dummy_path}\t\tAdd a path to an upscale model directory.") + print(f"\t--vae-dir {dummy_path}\t\t\tAdd a path to a vae directory.") print("\t--dont-upcast-attention\t\tDisable upcasting of attention \n\t\t\t\t\tcan boost speed but increase the chances of black images.\n") print("\t--use-split-cross-attention\tUse the split cross attention optimization instead of the sub-quadratic one.\n\t\t\t\t\tIgnored when xformers is used.") @@ -74,6 +74,18 @@ def cleanup_temp(): if os.path.exists(temp_dir): shutil.rmtree(temp_dir, ignore_errors=True) +def add_model_search_paths(): + if '--extra-model-paths-config' in sys.argv: + indices = [(i + 1) for i in range(len(sys.argv) - 1) if sys.argv[i] == '--extra-model-paths-config'] + for i in indices: + load_extra_path_config(sys.argv[i]) + + for model_tag, search_location in folder_paths.model_search_locations.items(): + model_path_arg = f"--{model_tag.replace('_', '-')}-dir" + extra_paths = [sys.argv[i + 1] for i, e in enumerate(sys.argv[:-1]) if e == model_path_arg] + for extra_path in extra_paths: + folder_paths.add_model_search_path(model_tag, extra_path) + def load_extra_path_config(yaml_path): with open(yaml_path, 'r') as stream: config = yaml.safe_load(stream) @@ -92,7 +104,7 @@ def load_extra_path_config(yaml_path): if base_path is not None: full_path = os.path.join(base_path, full_path) print("Adding extra search path", x, full_path) - folder_paths.add_model_folder_path(x, full_path) + folder_paths.add_model_search_path(x, full_path) if __name__ == "__main__": cleanup_temp() @@ -118,10 +130,7 @@ if __name__ == "__main__": if os.path.isfile(extra_model_paths_config_path): load_extra_path_config(extra_model_paths_config_path) - if '--extra-model-paths-config' in sys.argv: - indices = [(i + 1) for i in range(len(sys.argv) - 1) if sys.argv[i] == '--extra-model-paths-config'] - for i in indices: - load_extra_path_config(sys.argv[i]) + add_model_search_paths() port = 8188 try: