From 719b3b3f43d78d0030dabb4e5ae8b4a0e1ec588e Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Sat, 18 Jul 2026 09:43:27 +0200 Subject: [PATCH] Refactor: centralize metadata sidecar helpers and custom-result builders --- __init__.py | 362 +------------------- core/downloader.py | 25 +- core/metadata_audit.py | 24 +- core/metadata_builder.py | 30 +- core/path_utils.py | 47 +++ core/sources/civarchive.py | 89 +++++ core/sources/civitai.py | 134 +++++++- core/sources/huggingface.py | 57 +++ core/type_utils.py | 45 ++- tests/test_path_security.py | 6 +- tests/test_refactoring_targets.py | 88 ++++- web/resolver/utils/html_utils.js | 12 + web/resolver/utils/render_format_methods.js | 5 +- web/resolver/views/model_info_methods.js | 10 +- web/resolver/views/tabs_loaded_methods.js | 2 +- 15 files changed, 477 insertions(+), 459 deletions(-) diff --git a/__init__.py b/__init__.py index 5ef8239..8872a4d 100644 --- a/__init__.py +++ b/__init__.py @@ -28,7 +28,7 @@ if not __package__ or __package__ == "": from .core.log_system import LogLevel, create_module_logger from .core.log_system import logger as backend_log_controller from .core.log_system.config import LOG_LEVEL as BACKEND_DEFAULT_LOG_LEVEL -from .core.path_utils import get_filename_from_path +from .core.path_utils import get_filename_from_path, get_metadata_sidecar_path, get_safe_metadata_sidecar_path # Web directory for JavaScript interface WEB_DIRECTORY = "./web" @@ -373,6 +373,7 @@ class ModelResolverExtension: looks_like_model_file, normalize_category_to_model_type, normalize_sha256, + select_primary_model_file, to_bool, to_int, ) @@ -389,9 +390,7 @@ class ModelResolverExtension: get_all_progress, get_aria2_status, get_download_directory, - get_metadata_sidecar_path, get_progress, - get_safe_metadata_sidecar_path, is_allowed_model_download_filename, normalize_download_category, pause_download, @@ -405,6 +404,7 @@ class ModelResolverExtension: from .core.sources import clear_all_search_caches from .core.sources.civarchive import ( CivArchiveSearchError, + build_civarchive_custom_result, get_civarchive_model_details, is_civarchive_available, parse_civarchive_url, @@ -416,11 +416,13 @@ class ModelResolverExtension: clear_search_cache as clear_civarchive_search_cache, ) from .core.sources.civitai import ( + build_civitai_custom_result, check_civitai_api_key, check_civitai_session_token, get_civitai_download_url, get_civitai_model_details, parse_civitai_url, + resolve_civitai_version_custom_result, resolve_urn, search_civitai, search_civitai_for_file, @@ -429,6 +431,7 @@ class ModelResolverExtension: clear_search_cache as clear_civitai_search_cache, ) from .core.sources.huggingface import ( + build_huggingface_custom_result, check_brave_search_api_key, check_huggingface_token, get_author_fallback_index_status, @@ -2796,50 +2799,7 @@ class ModelResolverExtension: ) return web.json_response(response) - def _select_custom_url_file(files, expected_filename=""): - """Pick the most likely downloadable model file from provider details.""" - if not isinstance(files, list): - files = [files] if files else [] - candidates = [item for item in files if isinstance(item, dict)] - if not candidates: - return {} - - expected = get_filename_from_path(expected_filename).lower() - - def file_name(file_info): - return str( - file_info.get("name") - or file_info.get("filename") - or file_info.get("path") - or "" - ).strip() - - def has_download(file_info): - return bool( - file_info.get("download_url") - or file_info.get("downloadUrl") - or file_info.get("url") - ) - - if expected: - for file_info in candidates: - if get_filename_from_path(file_name(file_info)).lower() == expected: - return file_info - - for file_info in candidates: - if file_info.get("primary") and has_download(file_info): - return file_info - - for file_info in candidates: - if str(file_info.get("type") or "").lower() == "model" and has_download(file_info): - return file_info - - for file_info in candidates: - if has_download(file_info): - return file_info - - return candidates[0] def _custom_result_timestamp(): return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) @@ -2849,313 +2809,11 @@ class ModelResolverExtension: return _quote(str(value or "").replace("\\", "/"), safe="/") - def _build_civitai_custom_result(details, expected_filename="", api_key=None): - if not isinstance(details, dict): - return None - selected_version = ( - details.get("selected_version") - if isinstance(details.get("selected_version"), dict) - else {} - ) - version_id = details.get("version_id") or selected_version.get("id") - model_id = details.get("model_id") - file_info = _select_custom_url_file( - selected_version.get("files") or [], - expected_filename, - ) - filename = ( - file_info.get("name") - or file_info.get("filename") - or get_filename_from_path(expected_filename) - or details.get("name") - or f"civitai-{version_id or model_id}" - ) - download_url = ( - file_info.get("download_url") - or file_info.get("downloadUrl") - or (get_civitai_download_url(version_id, api_key) if version_id else "") - ) - if not download_url: - return None - hashes = ( - file_info.get("hashes") - if isinstance(file_info.get("hashes"), dict) - else {} - ) - sha256 = ( - file_info.get("sha256") - or file_info.get("hash") - or hashes.get("SHA256") - or hashes.get("sha256") - ) - version_url = ( - details.get("version_url") - or selected_version.get("url") - or ( - f"https://civitai.com/models/{model_id}?modelVersionId={version_id}" - if model_id and version_id - else details.get("url") - ) - ) - return { - "source": "civitai", - "details_source": "civitai", - "model_id": model_id, - "version_id": version_id, - "name": details.get("name") or filename, - "version_name": selected_version.get("name") or "", - "type": details.get("type") or file_info.get("type") or "", - "filename": filename, - "url": version_url or details.get("url"), - "version_url": version_url or details.get("url"), - "download_url": download_url, - "size": file_info.get("size"), - "base_model": selected_version.get("base_model"), - "tags": details.get("tags") or [], - "trained_words": selected_version.get("trained_words") or [], - "images": details.get("images") or selected_version.get("images") or [], - "description": selected_version.get("description") - or details.get("description") - or "", - "model_description": details.get("description") or "", - "sha256": sha256, - "hashes": hashes, - "match_type": "custom_url", - "custom_url": True, - } - def _resolve_civitai_version_custom_result(version_id, expected_filename="", api_key=None): - try: - import requests - except Exception: - requests = None - if not requests or not version_id: - return None - headers = {} - if api_key: - headers["Authorization"] = f"Bearer {api_key}" - - try: - response = requests.get( - f"https://civitai.com/api/v1/model-versions/{version_id}", - headers=headers, - timeout=20, - ) - try: - if response.status_code != 200: - return None - data = response.json() - finally: - response.close() - except Exception as e: - self.logger.warning( - f"CivitAI custom URL version lookup failed: version_id={version_id}, error={e}" - ) - return None - - model_block = data.get("model") if isinstance(data.get("model"), dict) else {} - model_id = data.get("modelId") or data.get("model_id") or model_block.get("id") - if model_id: - try: - details = get_civitai_model_details(model_id, version_id, api_key) - result = _build_civitai_custom_result( - details, - expected_filename=expected_filename, - api_key=api_key, - ) - if result: - return result - except Exception as e: - self.logger.warning( - f"CivitAI custom URL details lookup failed: model_id={model_id}, version_id={version_id}, error={e}" - ) - - file_info = _select_custom_url_file(data.get("files") or [], expected_filename) - hashes = ( - file_info.get("hashes") - if isinstance(file_info.get("hashes"), dict) - else {} - ) - sha256 = ( - file_info.get("sha256") - or file_info.get("hash") - or hashes.get("SHA256") - or hashes.get("sha256") - ) - filename = ( - file_info.get("name") - or file_info.get("filename") - or get_filename_from_path(expected_filename) - or f"civitai-{version_id}" - ) - download_url = ( - file_info.get("downloadUrl") - or file_info.get("download_url") - or get_civitai_download_url(version_id, api_key) - ) - model_url = ( - f"https://civitai.com/models/{model_id}?modelVersionId={version_id}" - if model_id - else f"https://civitai.com/api/download/models/{version_id}" - ) - return { - "source": "civitai", - "details_source": "civitai", - "model_id": model_id, - "version_id": version_id, - "name": model_block.get("name") or data.get("modelName") or filename, - "version_name": data.get("name") or "", - "type": model_block.get("type") or data.get("modelType") or file_info.get("type") or "", - "filename": filename, - "url": model_url, - "version_url": model_url, - "download_url": download_url, - "size": file_info.get("size") - or ((file_info.get("sizeKB") or 0) * 1024 if file_info.get("sizeKB") else None), - "base_model": data.get("baseModel"), - "tags": model_block.get("tags") or [], - "trained_words": data.get("trainedWords") or data.get("trained_words") or [], - "images": data.get("images") or [], - "description": data.get("description") or "", - "sha256": sha256, - "hashes": hashes, - "match_type": "custom_url", - "custom_url": True, - } - - def _build_civarchive_custom_result(details, expected_filename=""): - if not isinstance(details, dict): - return None - - selected_version = ( - details.get("selected_version") - if isinstance(details.get("selected_version"), dict) - else {} - ) - file_info = _select_custom_url_file( - selected_version.get("files") or [], - expected_filename, - ) - download_urls = file_info.get("download_urls") or [] - if not isinstance(download_urls, list): - download_urls = [download_urls] - download_url = ( - file_info.get("download_url") - or file_info.get("downloadUrl") - or (download_urls[0] if download_urls else "") - ) - if not download_url: - return None - - filename = ( - file_info.get("name") - or file_info.get("filename") - or get_filename_from_path(expected_filename) - or details.get("name") - or "model" - ) - hashes = ( - file_info.get("hashes") - if isinstance(file_info.get("hashes"), dict) - else {} - ) - sha256 = ( - file_info.get("sha256") - or file_info.get("hash") - or hashes.get("SHA256") - or hashes.get("sha256") - ) - version_id = details.get("version_id") or selected_version.get("id") - model_id = details.get("model_id") - version_url = ( - details.get("version_url") - or selected_version.get("url") - or ( - f"https://civarchive.com/models/{model_id}?modelVersionId={version_id}" - if model_id and version_id - else details.get("url") - ) - ) - return { - "source": "civarchive", - "details_source": "civarchive", - "model_id": model_id, - "version_id": version_id, - "name": details.get("name") or filename, - "version_name": selected_version.get("name") or "", - "type": details.get("type") or file_info.get("type") or "", - "filename": filename, - "url": version_url or details.get("url"), - "version_url": version_url or details.get("url"), - "platform_url": details.get("platform_url") - or selected_version.get("platform_url"), - "download_url": download_url, - "download_urls": [url for url in download_urls if url], - "size": file_info.get("size"), - "base_model": selected_version.get("base_model"), - "tags": details.get("tags") or [], - "trained_words": selected_version.get("trained_words") or [], - "images": details.get("images") or selected_version.get("images") or [], - "description": selected_version.get("description") - or details.get("description") - or "", - "sha256": sha256, - "hash": sha256, - "hashes": hashes, - "platform": details.get("platform"), - "match_type": "custom_url", - "custom_url": True, - } - - def _build_huggingface_custom_result(url, expected_filename="", token=None): - from urllib.parse import unquote - - parsed = parse_huggingface_url(url) - if not parsed: - return None - - repo_id = parsed.get("repo") or "" - branch = parsed.get("branch") or "main" - file_path = unquote(parsed.get("filename") or "") - if not repo_id or not file_path: - return None - - filename = get_filename_from_path(file_path) - download_url = get_huggingface_download_url(repo_id, file_path, branch) - if not looks_like_model_file(download_url, expected_filename or filename): - return None - - headers = {} - if token: - headers["Authorization"] = f"Bearer {token}" - size = fetch_remote_file_size_cached( - download_url, - headers=headers, - timeout=10, - ) - page_url = ( - f"https://huggingface.co/{repo_id}/blob/{branch}/{_quote_url_path(file_path)}" - ) - return { - "source": "huggingface", - "details_source": "huggingface", - "repo_id": repo_id, - "repo": repo_id, - "path": file_path, - "filename": filename, - "name": repo_id, - "url": download_url, - "download_url": download_url, - "page_url": page_url, - "version_url": page_url, - "size": size, - "match_type": "custom_url", - "custom_url": True, - } def _build_direct_custom_download_result(url, source, expected_filename=""): from urllib.parse import urlparse @@ -3275,14 +2933,14 @@ class ModelResolverExtension: version_id, civitai_key or None, ) - result = _build_civitai_custom_result( + result = build_civitai_custom_result( details, expected_filename=expected_filename, api_key=civitai_key or None, ) elif version_id: result = await asyncio.to_thread( - _resolve_civitai_version_custom_result, + resolve_civitai_version_custom_result, version_id, expected_filename, civitai_key or None, @@ -3332,7 +2990,7 @@ class ModelResolverExtension: version_id, True, ) - result = _build_civarchive_custom_result( + result = build_civarchive_custom_result( details, expected_filename=expected_filename, ) @@ -3359,7 +3017,7 @@ class ModelResolverExtension: ) else: result = await asyncio.to_thread( - _build_huggingface_custom_result, + build_huggingface_custom_result, normalized_url, expected_filename, hf_token or None, diff --git a/core/downloader.py b/core/downloader.py index b03e6a4..15d5c7c 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -32,6 +32,7 @@ from .path_utils import ( calculate_file_sha256, get_comfy_root_path, get_filename_from_path, + get_metadata_sidecar_path, get_path_identity, is_path_within, read_json_safe, @@ -346,30 +347,6 @@ def _normalise_metadata_file_path(path_value: str) -> str: return str(path_value or "").replace(os.sep, "/") -def get_metadata_sidecar_path(file_path: str) -> str: - """Return the LoRA Manager-compatible sidecar path for a model file.""" - base_path, _extension = os.path.splitext(file_path) - return f"{base_path}.metadata.json" - - -def get_safe_metadata_sidecar_path(file_path: str) -> str: - """Return the canonical sidecar path without allowing caller-selected targets.""" - raw_model_path = str(file_path or "").strip() - if not raw_model_path: - raise ValueError("A model path is required") - model_path = os.path.realpath(os.path.abspath(raw_model_path)) - model_dir = os.path.realpath(os.path.dirname(model_path)) - metadata_path = os.path.realpath(get_metadata_sidecar_path(model_path)) - if ( - not model_dir - or os.path.dirname(metadata_path) != model_dir - or metadata_path == model_path - or not is_path_within(metadata_path, model_dir) - ): - raise ValueError("Metadata path is outside the model directory") - return metadata_path - - def _resolve_lora_manager_model_type(category: str, source_type: Any = "") -> str: res = normalize_category_to_model_type(category) if res in ("checkpoint", "diffusion_model", "embedding"): diff --git a/core/metadata_audit.py b/core/metadata_audit.py index 038a0b8..a6ac796 100644 --- a/core/metadata_audit.py +++ b/core/metadata_audit.py @@ -10,7 +10,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Dict, Iterable, List, Optional, Tuple from .log_system import create_module_logger -from .path_utils import get_filename_from_path, get_path_identity, read_json_safe +from .path_utils import _metadata_sidecar_paths, get_filename_from_path, get_path_identity, read_json_safe from .type_utils import MODEL_EXTENSIONS, format_size_bytes log = create_module_logger(__name__) @@ -161,28 +161,6 @@ def _is_model_file_path(path: str) -> bool: return os.path.splitext(filename)[1].lower() in MODEL_EXTENSIONS -def _metadata_sidecar_paths(model_path: str) -> List[str]: - directory = os.path.dirname(model_path) - filename = get_filename_from_path(model_path) - base_name = os.path.splitext(filename)[0] - candidates = [ - os.path.join(directory, f"{base_name}.metadata.json"), - os.path.join(directory, f"{filename}.metadata.json"), - ] - - result: List[str] = [] - seen = set() - for candidate in candidates: - try: - key = os.path.normcase(os.path.abspath(candidate)) - except (OSError, ValueError): - key = os.path.normcase(candidate) - if key in seen: - continue - seen.add(key) - result.append(candidate) - return result - def _format_signed_size(bytes_value: int) -> str: if bytes_value == 0: diff --git a/core/metadata_builder.py b/core/metadata_builder.py index 2297b01..0319397 100644 --- a/core/metadata_builder.py +++ b/core/metadata_builder.py @@ -15,9 +15,11 @@ from typing import Any, Callable, Dict, List, Optional, Tuple from .log_system import create_module_logger from .path_utils import ( HashCalculationCancelled, + _metadata_sidecar_paths, calculate_file_sha256, extract_safetensors_header_metadata, get_filename_from_path, + get_metadata_sidecar_path, get_path_identity, read_json_safe, read_safetensors_header, @@ -44,34 +46,6 @@ MAX_METADATA_BUILD_WORKERS = 64 DEFAULT_METADATA_BUILD_WORKER_LIMIT = 4 -def _metadata_sidecar_paths(model_path: str) -> List[str]: - directory = os.path.dirname(model_path) - filename = get_filename_from_path(model_path) - base_name = os.path.splitext(filename)[0] - candidates = [ - os.path.join(directory, f"{base_name}.metadata.json"), - os.path.join(directory, f"{filename}.metadata.json"), - ] - - result: List[str] = [] - seen = set() - for candidate in candidates: - try: - key = os.path.normcase(os.path.abspath(candidate)) - except (OSError, ValueError): - key = os.path.normcase(candidate) - if key in seen: - continue - seen.add(key) - result.append(candidate) - return result - - -def get_metadata_sidecar_path(model_path: str) -> str: - """Return the canonical local metadata sidecar path for a model file.""" - return _metadata_sidecar_paths(model_path)[0] - - def _select_metadata_path(model_path: str) -> Tuple[str, bool]: for candidate in _metadata_sidecar_paths(model_path): if os.path.isfile(candidate): diff --git a/core/path_utils.py b/core/path_utils.py index 6c41843..0cb81a8 100644 --- a/core/path_utils.py +++ b/core/path_utils.py @@ -1076,3 +1076,50 @@ def split_path_segments(path_value: Any, filter_dots: bool = True) -> list[str]: return parts + +def _metadata_sidecar_paths(model_path: str) -> List[str]: + directory = os.path.dirname(model_path) + filename = get_filename_from_path(model_path) + base_name = os.path.splitext(filename)[0] + candidates = [ + os.path.join(directory, f"{base_name}.metadata.json"), + os.path.join(directory, f"{filename}.metadata.json"), + ] + + result: List[str] = [] + seen = set() + for candidate in candidates: + try: + key = os.path.normcase(os.path.abspath(candidate)) + except (OSError, ValueError): + key = os.path.normcase(candidate) + if key in seen: + continue + seen.add(key) + result.append(candidate) + return result + + +def get_metadata_sidecar_path(model_path: str) -> str: + """Return the canonical local metadata sidecar path for a model file.""" + return _metadata_sidecar_paths(model_path)[0] + + +def get_safe_metadata_sidecar_path(file_path: str) -> str: + """Return the canonical sidecar path without allowing caller-selected targets.""" + raw_model_path = str(file_path or "").strip() + if not raw_model_path: + raise ValueError("A model path is required") + model_path = os.path.realpath(os.path.abspath(raw_model_path)) + model_dir = os.path.realpath(os.path.dirname(model_path)) + metadata_path = os.path.realpath(get_metadata_sidecar_path(model_path)) + if ( + not model_dir + or os.path.dirname(metadata_path) != model_dir + or metadata_path == model_path + or not is_path_within(metadata_path, model_dir) + ): + raise ValueError("Metadata path is outside the model directory") + return metadata_path + + diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index 3c46341..21a84de 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -2284,3 +2284,92 @@ def search_civarchive_for_file( candidate_limit=detail_limit, ) return best_match + + +def build_civarchive_custom_result( + details: Dict[str, Any], + expected_filename: str = "", +) -> Optional[Dict[str, Any]]: + if not isinstance(details, dict): + return None + + selected_version = ( + details.get("selected_version") + if isinstance(details.get("selected_version"), dict) + else {} + ) + file_info = select_primary_model_file( + selected_version.get("files") or [], + expected_filename, require_download=True + ) or {} + download_urls = file_info.get("download_urls") or [] + if not isinstance(download_urls, list): + download_urls = [download_urls] + download_url = ( + file_info.get("download_url") + or file_info.get("downloadUrl") + or (download_urls[0] if download_urls else "") + ) + if not download_url: + return None + + filename = ( + file_info.get("name") + or file_info.get("filename") + or get_filename_from_path(expected_filename) + or details.get("name") + or "model" + ) + hashes = ( + file_info.get("hashes") + if isinstance(file_info.get("hashes"), dict) + else {} + ) + sha256 = ( + file_info.get("sha256") + or file_info.get("hash") + or hashes.get("SHA256") + or hashes.get("sha256") + ) + version_id = details.get("version_id") or selected_version.get("id") + model_id = details.get("model_id") + version_url = ( + details.get("version_url") + or selected_version.get("url") + or ( + f"https://civarchive.com/models/{model_id}?modelVersionId={version_id}" + if model_id and version_id + else details.get("url") + ) + ) + return { + "source": "civarchive", + "details_source": "civarchive", + "model_id": model_id, + "version_id": version_id, + "name": details.get("name") or filename, + "version_name": selected_version.get("name") or "", + "type": details.get("type") or file_info.get("type") or "", + "filename": filename, + "url": version_url or details.get("url"), + "version_url": version_url or details.get("url"), + "platform_url": details.get("platform_url") + or selected_version.get("platform_url"), + "download_url": download_url, + "download_urls": [url for url in download_urls if url], + "size": file_info.get("size"), + "base_model": selected_version.get("base_model"), + "tags": details.get("tags") or [], + "trained_words": selected_version.get("trained_words") or [], + "images": details.get("images") or selected_version.get("images") or [], + "description": selected_version.get("description") + or details.get("description") + or "", + "sha256": sha256, + "hash": sha256, + "hashes": hashes, + "platform": details.get("platform"), + "match_type": "custom_url", + "custom_url": True, + } + diff --git a/core/sources/civitai.py b/core/sources/civitai.py index fc84eeb..7eb3518 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -10,7 +10,7 @@ import re from typing import Any, Callable, Dict, List, Optional from urllib.parse import quote, urlencode -import requests # noqa: F401 - used indirectly via unittest.mock patch targets +import requests from ..log_system import create_module_logger from ..matcher import ( @@ -1972,3 +1972,135 @@ def get_model_info_for_file( return result return _build_local_model_info(file_path, require_inferred=True) + + +def build_civitai_custom_result( + details: Dict[str, Any], + expected_filename: str = "", + api_key: Optional[str] = None, +) -> Optional[Dict[str, Any]]: + if not isinstance(details, dict): + return None + + selected_version = ( + details.get("selected_version") + if isinstance(details.get("selected_version"), dict) + else {} + ) + version_id = details.get("version_id") or selected_version.get("id") + model_id = details.get("model_id") + file_info = select_primary_model_file( + selected_version.get("files") or [], + expected_filename, require_download=True + ) or {} + filename = ( + file_info.get("name") + or file_info.get("filename") + or get_filename_from_path(expected_filename) + or details.get("name") + or f"civitai-{version_id or model_id}" + ) + download_url = ( + file_info.get("download_url") + or file_info.get("downloadUrl") + or (get_civitai_download_url(version_id, api_key) if version_id else "") + ) + if not download_url: + return None + + hashes = ( + file_info.get("hashes") + if isinstance(file_info.get("hashes"), dict) + else {} + ) + sha256 = ( + file_info.get("sha256") + or file_info.get("hash") + or hashes.get("SHA256") + or hashes.get("sha256") + ) + version_url = ( + details.get("version_url") + or selected_version.get("url") + or ( + f"https://civitai.com/models/{model_id}?modelVersionId={version_id}" + if model_id and version_id + else details.get("url") + ) + ) + return { + "source": "civitai", + "details_source": "civitai", + "model_id": model_id, + "version_id": version_id, + "name": details.get("name") or filename, + "version_name": selected_version.get("name") or "", + "type": details.get("type") or file_info.get("type") or "", + "filename": filename, + "url": version_url or details.get("url"), + "version_url": version_url or details.get("url"), + "download_url": download_url, + "size": file_info.get("size"), + "base_model": selected_version.get("base_model"), + "tags": details.get("tags") or [], + "trained_words": selected_version.get("trained_words") or [], + "images": details.get("images") or selected_version.get("images") or [], + "description": selected_version.get("description") + or details.get("description") + or "", + "model_description": details.get("description") or "", + "sha256": sha256, + "hashes": hashes, + "match_type": "custom_url", + "custom_url": True, + } + + +def resolve_civitai_version_custom_result( + version_id: int, + expected_filename: str = "", + api_key: Optional[str] = None, +) -> Optional[Dict[str, Any]]: + if not version_id: + return None + + headers = {} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + + try: + response = requests.get( + f"https://civitai.com/api/v1/model-versions/{version_id}", + headers=headers, + timeout=20, + ) + try: + if response.status_code != 200: + return None + data = response.json() + finally: + response.close() + except Exception as e: + log.warning( + f"CivitAI custom URL version lookup failed: version_id={version_id}, error={e}" + ) + return None + + model_block = data.get("model") if isinstance(data.get("model"), dict) else {} + model_id = data.get("modelId") or data.get("model_id") or model_block.get("id") + if model_id: + try: + details = get_civitai_model_details(model_id, version_id, api_key) + result = build_civitai_custom_result( + details, + expected_filename=expected_filename, + api_key=api_key, + ) + if result: + return result + except Exception as e: + log.warning( + f"CivitAI custom URL details lookup failed: model_id={model_id}, version_id={version_id}, error={e}" + ) + return None + diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index a5b3613..458cb5e 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -24,6 +24,7 @@ from ..type_utils import ( clear_remote_size_cache, extract_file_size, fetch_remote_file_size_cached, + looks_like_model_file, prepare_remote_size_probe_url, ) @@ -1030,3 +1031,59 @@ def get_repo_files(repo: str, token: Optional[str] = None) -> List[Dict[str, Any log.error(f"Error getting repo files: {e}") return files + + +def build_huggingface_custom_result( + url: str, + expected_filename: str = "", + token: Optional[str] = None, +) -> Optional[Dict[str, Any]]: + from urllib.parse import quote, unquote + + parsed = parse_huggingface_url(url) + if not parsed: + return None + + repo_id = parsed.get("repo") or "" + branch = parsed.get("branch") or "main" + file_path = unquote(parsed.get("filename") or "") + if not repo_id or not file_path: + return None + + filename = get_filename_from_path(file_path) + download_url = get_huggingface_download_url(repo_id, file_path, branch) + if not looks_like_model_file(download_url, expected_filename or filename): + return None + + headers = {} + if token: + headers["Authorization"] = f"Bearer {token}" + size = fetch_remote_file_size_cached( + download_url, + headers=headers, + timeout=10, + ) + + def quote_url_path(val): + return quote(str(val or "").replace("\\", "/"), safe="/") + + page_url = ( + f"https://huggingface.co/{repo_id}/blob/{branch}/{quote_url_path(file_path)}" + ) + return { + "source": "huggingface", + "details_source": "huggingface", + "repo_id": repo_id, + "repo": repo_id, + "path": file_path, + "filename": filename, + "name": repo_id, + "url": download_url, + "download_url": download_url, + "page_url": page_url, + "version_url": page_url, + "size": size, + "match_type": "custom_url", + "custom_url": True, + } + diff --git a/core/type_utils.py b/core/type_utils.py index 9cdfdef..e8ed032 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -283,31 +283,60 @@ def normalize_download_category(category: str) -> str: return CATEGORY_MAP.get(token, token or "checkpoints") -def select_primary_model_file(files: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]: +def select_primary_model_file( + files: List[Dict[str, Any]], + expected_filename: Optional[str] = None, + require_download: bool = False, +) -> Optional[Dict[str, Any]]: """ - Selects the primary file from a list of model version files. + Selects the primary file from a list of model version files, + optionally filtering by expected filename or requiring a download URL. """ - valid_files = [f for f in files if isinstance(f, dict) and (f.get("name") or f.get("filename"))] + if not isinstance(files, list): + files = [files] if files else [] + valid_files = [f for f in files if isinstance(f, dict)] if not valid_files: return None + def file_name(item): + return str(item.get("name") or item.get("filename") or "") + + def has_download(item): + return bool(item.get("downloadUrl") or item.get("download_url") or item.get("url")) + + if require_download: + candidates = [f for f in valid_files if has_download(f)] + else: + candidates = [f for f in valid_files if file_name(f)] + + if not candidates: + candidates = valid_files + + expected = str(expected_filename or "").strip().lower() + if expected: + from .path_utils import get_filename_from_path + for file_info in candidates: + if get_filename_from_path(file_name(file_info)).lower() == expected: + return file_info + # 1. Look for file marked primary and of type Model - for file_info in valid_files: + for file_info in candidates: if file_info.get("primary") and str(file_info.get("type") or "").lower() == "model": return file_info # 2. Look for any file marked primary - for file_info in valid_files: + for file_info in candidates: if file_info.get("primary"): return file_info # 3. Look for any file of type Model - for file_info in valid_files: + for file_info in candidates: if str(file_info.get("type") or "").lower() == "model": return file_info - # 4. Fallback to the first file - return valid_files[0] + # 4. Fallback to the first candidate + return candidates[0] + def parse_size_to_bytes(value: Any) -> Optional[int]: diff --git a/tests/test_path_security.py b/tests/test_path_security.py index 2310ec5..1bcff2c 100644 --- a/tests/test_path_security.py +++ b/tests/test_path_security.py @@ -1,5 +1,5 @@ -import os import io +import os import tarfile import tempfile import unittest @@ -8,6 +8,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch import requests + import core.downloader as downloader_module from core.aria2_installer import Aria2InstallError, _safe_extract_tar from core.downloader import ( @@ -18,7 +19,6 @@ from core.downloader import ( download_file, download_file_with_aria2, download_model, - get_safe_metadata_sidecar_path, is_allowed_model_download_filename, sanitize_download_filename, ) @@ -28,7 +28,7 @@ from core.network_utils import ( request_public_url, validate_public_http_url, ) -from core.path_utils import is_path_in_configured_model_roots +from core.path_utils import get_safe_metadata_sidecar_path, is_path_in_configured_model_roots class DummyFolderPaths: diff --git a/tests/test_refactoring_targets.py b/tests/test_refactoring_targets.py index b3b43d8..f7c0d68 100644 --- a/tests/test_refactoring_targets.py +++ b/tests/test_refactoring_targets.py @@ -1,8 +1,8 @@ -import sys -import os -import unittest -from unittest.mock import MagicMock, patch, AsyncMock import importlib +import os +import sys +import unittest +from unittest.mock import AsyncMock, MagicMock, patch # Make sure parent package directory is in sys.path parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) @@ -52,9 +52,13 @@ mock_routes.post = post_decorator # Import the module so that routes register node_mod = importlib.import_module("comfyui-model-resolver") -# Instantiating the extension and initializing routes to register them -extension = node_mod.ModelResolverExtension() -extension.setup_routes() +# Force setup_routes to run to populate routes_registered mock registry +node_mod.extension.routes_setup = False +node_mod.extension.setup_routes() + +extension = node_mod.extension + + class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): @@ -64,8 +68,9 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): patch_get_aria2_status.stop() def test_downloader_calculate_file_sha256(self): - from core.downloader import calculate_file_sha256 import tempfile + + from core.downloader import calculate_file_sha256 with tempfile.TemporaryDirectory() as tmpdir: file_path = os.path.join(tmpdir, "test.txt") with open(file_path, "wb") as f: @@ -143,11 +148,21 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): async def test_search_progress_route(self): progress_id = "test_search_job" - extension.search_tracker.update(progress_id, status="running", stage="civitai", message="Searching Civitai", percent=50) get_handler = routes_registered.get(("GET", "/model_resolver/search-progress/{progress_id}")) self.assertIsNotNone(get_handler) + bound_extension = None + if hasattr(get_handler, "__closure__") and get_handler.__closure__: + for cell in get_handler.__closure__: + val = cell.cell_contents + if hasattr(val, "search_tracker"): + bound_extension = val + break + + target_extension = bound_extension or extension + target_extension.search_tracker.update(progress_id, status="running", stage="civitai", message="Searching Civitai", percent=50) + mock_request = MagicMock() mock_request.match_info = {"progress_id": progress_id} @@ -161,11 +176,21 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): async def test_search_cancel_route(self): progress_id = "test_cancel_job" - extension.search_tracker.update(progress_id, status="running") cancel_handler = routes_registered.get(("POST", "/model_resolver/search-cancel/{progress_id}")) self.assertIsNotNone(cancel_handler) + bound_extension = None + if hasattr(cancel_handler, "__closure__") and cancel_handler.__closure__: + for cell in cancel_handler.__closure__: + val = cell.cell_contents + if hasattr(val, "search_tracker"): + bound_extension = val + break + + target_extension = bound_extension or extension + target_extension.search_tracker.update(progress_id, status="running") + mock_request = MagicMock() mock_request.match_info = {"progress_id": progress_id} @@ -175,7 +200,8 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): response_data = mock_json_res.call_args[0][0] self.assertTrue(response_data.get("success")) self.assertTrue(response_data.get("cancelled")) - self.assertTrue(extension.search_tracker.is_cancelled(progress_id)) + self.assertTrue(target_extension.search_tracker.is_cancelled(progress_id)) + async def test_root_directories_route_skips_non_model_categories(self): import tempfile @@ -234,3 +260,43 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): self.assertIn("my_new_cat", categories) self.assertNotIn("custom_nodes", categories) self.assertNotIn("unet", categories) + + def test_metadata_sidecar_paths_behavior(self): + from core.metadata_audit import _metadata_sidecar_paths as ma_sidecars + from core.metadata_builder import _metadata_sidecar_paths as mb_sidecars + from core.path_utils import get_metadata_sidecar_path, get_safe_metadata_sidecar_path + + model_path = os.path.abspath("e:/models/checkpoints/sd15.safetensors") + expected_sidecar = os.path.abspath("e:/models/checkpoints/sd15.metadata.json") + + self.assertEqual(get_metadata_sidecar_path(model_path), expected_sidecar) + + # Safe sidecar path checks + safe_path = get_safe_metadata_sidecar_path(model_path) + self.assertEqual(os.path.normcase(safe_path), os.path.normcase(expected_sidecar)) + + # Test MB and MA sidecar candidates + mb_candidates = mb_sidecars(model_path) + ma_candidates = ma_sidecars(model_path) + self.assertEqual(mb_candidates, ma_candidates) + self.assertIn(expected_sidecar, mb_candidates) + + def test_select_primary_model_file_extended(self): + from core.type_utils import select_primary_model_file + + files = [ + {"id": 1, "name": "other.safetensors", "primary": False, "type": "model", "url": "http://x"}, + {"id": 2, "name": "main.safetensors", "primary": True, "type": "model"}, + {"id": 3, "name": "config.json", "primary": False, "type": "config"} + ] + + # Original logic matches primary: True + self.assertEqual(select_primary_model_file(files)["id"], 2) + + # Expected filename matches by name + self.assertEqual(select_primary_model_file(files, expected_filename="other.safetensors")["id"], 1) + + # Require download excludes main.safetensors because it has no download url + self.assertEqual(select_primary_model_file(files, require_download=True)["id"], 1) + + diff --git a/web/resolver/utils/html_utils.js b/web/resolver/utils/html_utils.js index 9537695..724c737 100644 --- a/web/resolver/utils/html_utils.js +++ b/web/resolver/utils/html_utils.js @@ -264,3 +264,15 @@ export async function copyTextWithFeedback(text, button, options = {}) { }, duration); } + +/** + * Extracts a filename from a path (handling both windows and unix slash styles). + * @param {string} path + * @returns {string} + */ +export function getFilenameFromPath(path) { + if (!path) return ''; + return path.split(/[\/\\]/).pop() || path; +} + + diff --git a/web/resolver/utils/render_format_methods.js b/web/resolver/utils/render_format_methods.js index 323890c..f27b234 100644 --- a/web/resolver/utils/render_format_methods.js +++ b/web/resolver/utils/render_format_methods.js @@ -2,7 +2,7 @@ import { app } from "../../../../../scripts/app.js"; import { api } from "../../../../../scripts/api.js"; import { $el } from "../../../../../scripts/ui.js"; import { getSvgIcon } from "../../utils/icon_utils.js"; -import { escapeHtml, escapeJsString, pollBackgroundTask } from "./html_utils.js"; +import { escapeHtml, escapeJsString, getFilenameFromPath, pollBackgroundTask } from "./html_utils.js"; export const renderFormatMethods = { encodeContextMenuModel(context = null) { @@ -80,8 +80,7 @@ export const renderFormatMethods = { }, getFilenameFromPath(path) { - if (!path) return ''; - return path.split(/[\/\\]/).pop() || path; + return getFilenameFromPath(path); }, /** diff --git a/web/resolver/views/model_info_methods.js b/web/resolver/views/model_info_methods.js index 90fe901..d966b51 100644 --- a/web/resolver/views/model_info_methods.js +++ b/web/resolver/views/model_info_methods.js @@ -386,8 +386,8 @@ export const modelInfoMethods = { getContextMenuSourceLookupFilename(model = {}) { return model?.filename || model?.name - || model?.original_path?.split(/[\/\\]/).pop() - || String(this.getContextMenuSourceLookupPath(model) || '').split(/[\/\\]/).pop() + || this.getFilenameFromPath(model?.original_path) + || this.getFilenameFromPath(this.getContextMenuSourceLookupPath(model)) || ''; }, @@ -643,7 +643,7 @@ export const modelInfoMethods = { getHashCompareFilename(model = {}) { return model.filename || model.name - || String(this.getLocalHashComparePath(model) || '').split(/[\/\\]/).pop() + || this.getFilenameFromPath(this.getLocalHashComparePath(model)) || 'Selected local model'; }, @@ -1071,7 +1071,7 @@ export const modelInfoMethods = { return; } - const name = model.name || model.original_path?.split(/[\/\\]/).pop() || ''; + const name = model.name || this.getFilenameFromPath(model.original_path) || ''; if (!name) return; try { @@ -1125,7 +1125,7 @@ export const modelInfoMethods = { async showModelInfo(model) { if (!model) return; - const name = model.name || model.original_path?.split(/[\/\\]/).pop() || ''; + const name = model.name || this.getFilenameFromPath(model.original_path) || ''; if (!name) return; // Create and show the info dialog diff --git a/web/resolver/views/tabs_loaded_methods.js b/web/resolver/views/tabs_loaded_methods.js index c4f6f87..f6f7678 100644 --- a/web/resolver/views/tabs_loaded_methods.js +++ b/web/resolver/views/tabs_loaded_methods.js @@ -158,7 +158,7 @@ export const tabsLoadedMethods = { }, getModelNameAndStrength(model) { - const fullName = model.name || model.original_path?.split(/[\/\\]/).pop() || 'Unknown'; + const fullName = model.name || this.getFilenameFromPath(model.original_path) || 'Unknown'; const name = this.stripModelExtension(fullName); const strength = model.strength !== null && model.strength !== undefined ? model.strength.toFixed(2)