From dada903b2bfe416b35a96a5f92c045758bf3016e Mon Sep 17 00:00:00 2001 From: Hayden <48267247+hayden-fr@users.noreply.github.com> Date: Thu, 28 Nov 2024 12:04:23 +0800 Subject: [PATCH] Feature scan extra folders (#65) * scan extra folders Other extension may be add models folder in folder_paths * Fix scanned non-model files Model file suffix specified --- __init__.py | 3 +-- py/config.py | 1 - py/services.py | 12 ++++++------ py/utils.py | 14 +++++++------- 4 files changed, 14 insertions(+), 16 deletions(-) diff --git a/__init__.py b/__init__.py index 2ca3a62..cf5eefc 100644 --- a/__init__.py +++ b/__init__.py @@ -23,7 +23,6 @@ if len(uninstalled_package) > 0: # Init config settings config.extension_uri = extension_uri -utils.resolve_model_base_paths() version = utils.get_current_version() utils.download_web_distribution(version) @@ -95,7 +94,7 @@ async def get_model_paths(request): """ Returns the base folders for models. """ - model_base_paths = config.model_base_paths + model_base_paths = utils.resolve_model_base_paths() return web.json_response({"success": True, "data": model_base_paths}) diff --git a/py/config.py b/py/config.py index 7f6a810..46471ce 100644 --- a/py/config.py +++ b/py/config.py @@ -1,7 +1,6 @@ extension_tag = "ComfyUI Model Manager" extension_uri: str = None -model_base_paths: dict[str, list[str]] = {} setting_key = { diff --git a/py/services.py b/py/services.py index e1e610d..d112548 100644 --- a/py/services.py +++ b/py/services.py @@ -10,14 +10,14 @@ from . import searcher def scan_models(request): result = [] - model_base_paths = config.model_base_paths + model_base_paths = utils.resolve_model_base_paths() for model_type in model_base_paths: folders, extensions = folder_paths.folder_names_and_paths[model_type] for path_index, base_path in enumerate(folders): files = utils.recursive_search_files(base_path, request) - models = folder_paths.filter_files_extensions(files, extensions) + models = folder_paths.filter_files_extensions(files, folder_paths.supported_pt_extensions) for fullname in models: fullname = utils.normalize_path(fullname) @@ -138,14 +138,14 @@ def fetch_model_info(model_page: str): async def download_model_info(scan_mode: str, request): utils.print_info(f"Download model info for {scan_mode}") - model_base_paths = config.model_base_paths + model_base_paths = utils.resolve_model_base_paths() for model_type in model_base_paths: folders, extensions = folder_paths.folder_names_and_paths[model_type] for path_index, base_path in enumerate(folders): files = utils.recursive_search_files(base_path, request) - models = folder_paths.filter_files_extensions(files, extensions) + models = folder_paths.filter_files_extensions(files, folder_paths.supported_pt_extensions) for fullname in models: fullname = utils.normalize_path(fullname) @@ -199,14 +199,14 @@ async def migrate_legacy_information(request): utils.print_info(f"Migrating legacy information...") - model_base_paths = config.model_base_paths + model_base_paths = utils.resolve_model_base_paths() for model_type in model_base_paths: folders, extensions = folder_paths.folder_names_and_paths[model_type] for path_index, base_path in enumerate(folders): files = utils.recursive_search_files(base_path, request) - models = folder_paths.filter_files_extensions(files, extensions) + models = folder_paths.filter_files_extensions(files, folder_paths.supported_pt_extensions) for fullname in models: fullname = utils.normalize_path(fullname) diff --git a/py/utils.py b/py/utils.py index b70e691..275e636 100644 --- a/py/utils.py +++ b/py/utils.py @@ -118,21 +118,21 @@ def download_web_distribution(version: str): def resolve_model_base_paths(): folders = list(folder_paths.folder_names_and_paths.keys()) - config.model_base_paths = {} + model_base_paths = {} + folder_black_list = ["configs", "custom_nodes"] for folder in folders: - if folder == "configs": - continue - if folder == "custom_nodes": + if folder in folder_black_list: continue folders = folder_paths.get_folder_paths(folder) - config.model_base_paths[folder] = [normalize_path(f) for f in folders] + model_base_paths[folder] = [normalize_path(f) for f in folders] + return model_base_paths def get_full_path(model_type: str, path_index: int, filename: str): """ Get the absolute path in the model type through string concatenation. """ - folders = config.model_base_paths.get(model_type, []) + folders = resolve_model_base_paths().get(model_type, []) if not path_index < len(folders): raise RuntimeError(f"PathIndex {path_index} is not in {model_type}") base_path = folders[path_index] @@ -144,7 +144,7 @@ def get_valid_full_path(model_type: str, path_index: int, filename: str): """ Like get_full_path but it will check whether the file is valid. """ - folders = config.model_base_paths.get(model_type, []) + folders = resolve_model_base_paths().get(model_type, []) if not path_index < len(folders): raise RuntimeError(f"PathIndex {path_index} is not in {model_type}") base_path = folders[path_index]