From 5de150053743cd50e2843606bfe8bfd20114f8b2 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Sat, 4 Jul 2026 00:04:22 +0200 Subject: [PATCH] refactor: consolidate inline mappings and precondition checks into shared helpers --- __init__.py | 83 +++---------------------------------- core/downloader.py | 12 ++---- core/scanner.py | 40 ++++++++++++++++++ core/sources/civitai.py | 26 +++++------- core/sources/huggingface.py | 26 +++++------- core/type_utils.py | 62 +++++++++++++++++++++++++-- 6 files changed, 129 insertions(+), 120 deletions(-) diff --git a/__init__.py b/__init__.py index 7f849ef..5dd5b6b 100644 --- a/__init__.py +++ b/__init__.py @@ -188,7 +188,7 @@ class ModelResolverExtension: search_local_matches, ) from .core.path_templates import infer_download_path_templates - from .core.scanner import get_model_files, invalidate_model_files_cache + from .core.scanner import get_model_files, invalidate_model_files_cache, find_local_file_path from .core.path_utils import ( get_filename_from_path, read_json_safe, @@ -204,6 +204,7 @@ class ModelResolverExtension: extract_sha256_from_metadata, fetch_remote_file_size_cached, looks_like_model_file, + normalize_category_to_model_type, ) from .core.settings import ( TEMPLATE_KEY_ALIASES, @@ -1515,33 +1516,8 @@ class ModelResolverExtension: file_path = resolved_path if resolved_path else None file_location = "" - if not file_path and category: - # Try to find the file in the model directories using folder_paths - try: - import folder_paths - - folder_type = normalize_download_category(category) - file_path = folder_paths.get_full_path( - folder_type, filename - ) - except Exception: - pass - - # If not found, try scanner - if not file_path: - try: - from .core.scanner import get_model_files - - available_models = get_model_files() - for m in available_models: - if ( - m.get("relative_path", "").endswith(filename) - or m.get("filename", "") == filename - ): - file_path = m.get("path") - break - except Exception: - pass + if not file_path: + file_path = find_local_file_path(filename, category) if file_path and _os.path.exists(file_path): file_location = _os.path.dirname(file_path).replace("\\", "/") @@ -1549,21 +1525,7 @@ class ModelResolverExtension: file_location += "/" def infer_model_type_from_category(value): - if not str(value or "").strip(): - return "" - normalized = normalize_download_category(value or "") - type_map = { - "checkpoints": "checkpoint", - "loras": "lora", - "vae": "vae", - "text_encoders": "text_encoder", - "clip": "clip", - "clip_vision": "clip_vision", - "controlnet": "controlnet", - "upscale_models": "upscale", - "diffusion_models": "diffusion_model", - } - return type_map.get(normalized, normalized.rstrip("s")) + return normalize_category_to_model_type(value) def build_info_response( result=None, @@ -2668,39 +2630,6 @@ class ModelResolverExtension: elif isinstance(result, dict): yield result - def get_result_sha256(result): - if not isinstance(result, dict): - return "" - hashes = ( - result.get("hashes") - if isinstance(result.get("hashes"), dict) - else {} - ) - file_info = ( - result.get("file_info") - if isinstance(result.get("file_info"), dict) - else {} - ) - file_hashes = ( - file_info.get("hashes") - if isinstance(file_info.get("hashes"), dict) - else {} - ) - for candidate in ( - result.get("sha256"), - result.get("hash"), - hashes.get("SHA256"), - hashes.get("sha256"), - file_info.get("sha256"), - file_info.get("hash"), - file_hashes.get("SHA256"), - file_hashes.get("sha256"), - ): - normalized = normalize_sha256(candidate) - if normalized: - return normalized - return "" - def collect_local_hash_matches(payload): matches = [] seen_match_paths = set() @@ -2715,7 +2644,7 @@ class ModelResolverExtension: ): for source_result in iter_result_items(payload.get(source_key)): raise_if_search_cancelled(source_key) - sha256 = get_result_sha256(source_result) + sha256 = extract_sha256_from_metadata(source_result) if not sha256: continue diff --git a/core/downloader.py b/core/downloader.py index e180037..9aeb783 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -25,7 +25,7 @@ log = create_module_logger(__name__) from .resolver import invalidate_local_hash_match_cache, normalize_sha256 from .scanner import invalidate_model_files_cache from .path_utils import is_path_within, get_path_identity, write_json_atomic, read_json_safe, get_comfy_root_path, calculate_file_sha256 as _calculate_file_sha256, get_filename_from_path -from .type_utils import as_dict, as_list, first_non_empty, format_size_bytes as format_bytes, DEFAULT_BROWSER_USER_AGENT, extract_response_file_size +from .type_utils import as_dict, as_list, first_non_empty, format_size_bytes as format_bytes, DEFAULT_BROWSER_USER_AGENT, extract_response_file_size, normalize_category_to_model_type try: import folder_paths @@ -259,13 +259,9 @@ def get_metadata_sidecar_path(file_path: str) -> str: def _resolve_lora_manager_model_type(category: str, source_type: Any = "") -> str: - category_key = normalize_download_category(category) - if category_key == "diffusion_models": - return "diffusion_model" - if category_key == "checkpoints": - return "checkpoint" - if category_key == "embeddings": - return "embedding" + res = normalize_category_to_model_type(category) + if res in ("checkpoint", "diffusion_model", "embedding"): + return res source_token = ( str(source_type or "") diff --git a/core/scanner.py b/core/scanner.py index 062692d..ede0594 100644 --- a/core/scanner.py +++ b/core/scanner.py @@ -306,3 +306,43 @@ def get_model_files(force_rescan: bool = False) -> List[Dict[str, str]]: _MODEL_FILES_CACHE = models _MODEL_FILES_CACHE_AT = now return models + + +def find_local_file_path(filename: str, category: Optional[str] = None) -> Optional[str]: + """ + Tries to find the absolute path of a model file locally using folder_paths or the local scanner. + """ + if not filename: + return None + + global folder_paths + if folder_paths is None: + try: + import folder_paths as fp + folder_paths = fp + except ImportError: + pass + + file_path = None + if category and folder_paths is not None: + try: + from .type_utils import normalize_download_category + folder_type = normalize_download_category(category) + file_path = folder_paths.get_full_path(folder_type, filename) + except Exception: + pass + + if not file_path: + try: + available_models = get_model_files() + for m in available_models: + if ( + m.get("relative_path", "").endswith(filename) + or m.get("filename", "") == filename + ): + file_path = m.get("path") + break + except Exception: + pass + + return file_path diff --git a/core/sources/civitai.py b/core/sources/civitai.py index b389659..d4ef32f 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -76,15 +76,12 @@ def build_civitai_session_cookie(session_token: Optional[str]) -> str: def check_civitai_session_token(session_token: Optional[str]) -> Dict[str, Any]: """Check whether a CivitAI browser session token is accepted by civitai.com.""" - token = (session_token or "").strip() - if not token: - return { - "success": False, - "valid": False, - "status": "missing", - "message": "Paste a CivitAI session token first.", - } + from ..type_utils import check_credential_preconditions + precheck = check_credential_preconditions(session_token, "CivitAI session token") + if precheck: + return precheck + token = (session_token or "").strip() headers = { "accept": "application/json", "Cookie": build_civitai_session_cookie(token), @@ -105,15 +102,12 @@ def check_civitai_session_token(session_token: Optional[str]) -> Dict[str, Any]: def check_civitai_api_key(api_key: Optional[str]) -> Dict[str, Any]: """Check whether a CivitAI API key is accepted by civitai.com.""" - key = (api_key or "").strip() - if not key: - return { - "success": False, - "valid": False, - "status": "missing", - "message": "Paste a CivitAI API key first.", - } + from ..type_utils import check_credential_preconditions + precheck = check_credential_preconditions(api_key, "CivitAI API key") + if precheck: + return precheck + key = (api_key or "").strip() headers = { "accept": "application/json", "Authorization": f"Bearer {key}", diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index 6d98ffb..7bb1f3a 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -51,15 +51,12 @@ _report_progress = get_progress_reporter("HuggingFace progress callback") def check_huggingface_token(token: Optional[str]) -> Dict[str, Any]: """Check whether a HuggingFace access token is accepted.""" - value = (token or "").strip() - if not value: - return { - "success": False, - "valid": False, - "status": "missing", - "message": "Paste a HuggingFace token first.", - } + from ..type_utils import check_credential_preconditions + precheck = check_credential_preconditions(token, "HuggingFace token") + if precheck: + return precheck + value = (token or "").strip() headers = {"Authorization": f"Bearer {value}"} def get_user(data): @@ -85,15 +82,12 @@ def check_huggingface_token(token: Optional[str]) -> Dict[str, Any]: def check_brave_search_api_key(api_key: Optional[str]) -> Dict[str, Any]: """Check whether a Brave Search API key is accepted.""" - key = (api_key or "").strip() - if not key: - return { - "success": False, - "valid": False, - "status": "missing", - "message": "Paste a Brave Search API key first.", - } + from ..type_utils import check_credential_preconditions + precheck = check_credential_preconditions(api_key, "Brave Search API key") + if precheck: + return precheck + key = (api_key or "").strip() headers = {"X-Subscription-Token": key} params = {"q": "test", "count": 1} diff --git a/core/type_utils.py b/core/type_utils.py index 67ec3fe..85cb594 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -743,12 +743,30 @@ def extract_sha256_from_metadata(metadata: Any) -> str: if normalized: return normalized + # Obsługa słownika file_info (np. z popular/model_list) + file_info = metadata.get("file_info") + if isinstance(file_info, dict): + for key in ("sha256", "hash", "SHA256", "Sha256"): + val = file_info.get(key) + if val: + normalized = normalize_sha256(val) + if normalized: + return normalized + f_hashes = file_info.get("hashes") + if isinstance(f_hashes, dict): + for key in ("SHA256", "sha256", "Sha256", "hash"): + val = f_hashes.get(key) + if val: + normalized = normalize_sha256(val) + if normalized: + return normalized + # Obsługa zagnieżdżonej listy plików w metadanych Civitai files = metadata.get("files") if isinstance(files, list): - for file_info in files: - if isinstance(file_info, dict): - f_hashes = file_info.get("hashes") + for file_info_item in files: + if isinstance(file_info_item, dict): + f_hashes = file_info_item.get("hashes") if isinstance(f_hashes, dict): for key in ("SHA256", "sha256", "Sha256", "hash"): val = f_hashes.get(key) @@ -878,5 +896,43 @@ def extract_file_size(file_info: Dict[str, Any]) -> Optional[int]: return None +def normalize_category_to_model_type(category: str) -> str: + """ + Mapuje nazwę kategorii pobierania/ComfyUI (np. 'checkpoints') na znormalizowany typ modelu (np. 'checkpoint'). + """ + if not category: + return "" + normalized = normalize_download_category(category) + type_map = { + "checkpoints": "checkpoint", + "loras": "lora", + "vae": "vae", + "text_encoders": "text_encoder", + "clip": "clip", + "clip_vision": "clip_vision", + "controlnet": "controlnet", + "upscale_models": "upscale", + "diffusion_models": "diffusion_model", + "embeddings": "embedding", + } + return type_map.get(normalized, normalized.rstrip("s")) + + +def check_credential_preconditions(value: Optional[str], token_name: str) -> Optional[Dict[str, Any]]: + """ + Sprawdza, czy poświadczenie (np. token, klucz) nie jest puste. + Zwraca standardowy słownik błędu, jeśli poświadczenie jest puste, w przeciwnym razie None. + """ + clean_val = (value or "").strip() + if not clean_val: + return { + "success": False, + "valid": False, + "status": "missing", + "message": f"Paste a {token_name} first.", + } + return None + +