From 38dc3e44744a606ab7de290491aaaeba37fc8d62 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Thu, 16 Jul 2026 10:52:06 +0200 Subject: [PATCH] refactor: centralize shared utilities, normalize imports, and configure project linting --- .gitignore | 2 + __init__.py | 225 +++++++++++++++------------ core/__init__.py | 5 +- core/aria2_installer.py | 1 - core/catalog_manager.py | 2 + core/downloader.py | 48 ++++-- core/log_system/__init__.py | 4 +- core/log_system/log_funcs.py | 7 +- core/log_system/logger.py | 15 +- core/matcher.py | 29 ++-- core/network_utils.py | 2 +- core/path_templates.py | 1 - core/path_utils.py | 27 ++-- core/progress.py | 4 +- core/resolver.py | 34 ++-- core/scanner.py | 6 +- core/settings.py | 4 +- core/sources/__init__.py | 68 ++++---- core/sources/civarchive.py | 50 +++--- core/sources/civitai.py | 60 +++---- core/sources/huggingface.py | 23 +-- core/sources/lora_manager_archive.py | 23 ++- core/sources/model_list.py | 17 +- core/sources/popular.py | 56 ++++--- core/type_utils.py | 24 +-- core/workflow_analyzer.py | 13 +- core/workflow_updater.py | 6 +- pyproject.toml | 46 ++++++ tests/test_huggingface.py | 3 +- 29 files changed, 463 insertions(+), 342 deletions(-) diff --git a/.gitignore b/.gitignore index fa6d669..3cdae95 100644 --- a/.gitignore +++ b/.gitignore @@ -36,3 +36,5 @@ tools/aria2/ # User settings – contains API keys and tokens, must NOT be committed model_resolver_settings.json + +node_modules/* \ No newline at end of file diff --git a/__init__.py b/__init__.py index 5b22da2..4700678 100644 --- a/__init__.py +++ b/__init__.py @@ -6,11 +6,11 @@ """ import asyncio -import threading -from typing import Dict, Any, Optional -import time -import sys import os +import sys +import threading +import time +from typing import Any, Dict, Optional if not __package__ or __package__ == "": this_dir = os.path.dirname(os.path.abspath(__file__)) @@ -25,7 +25,8 @@ if not __package__ or __package__ == "": sys.modules[package_name] = current_module if not hasattr(current_module, "__path__"): current_module.__path__ = [this_dir] -from .core.log_system import LogLevel, create_module_logger, logger as backend_log_controller +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 @@ -99,7 +100,7 @@ else: MODEL_RESOLVER_DEPENDENCY_NODE_TYPE: MODEL_RESOLVER_DEPENDENCY_NODE_DISPLAY_NAME, } - __all__ = ["WEB_DIRECTORY", "NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] + __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] class JobProgressTracker: @@ -140,7 +141,7 @@ class JobProgressTracker: now = time.time() with self.lock: current = self.progress.get(progress_id, {}) - + # If the job was cancelled, force it to remain cancelled. if progress_id in self.cancelled and payload.get("status") != "cancelled": payload["status"] = "cancelled" @@ -320,130 +321,146 @@ class ModelResolverExtension: # Import resolver modules try: - from .core.resolver import ( - analyze_and_find_matches, - apply_resolution, - get_local_model_hash_metadata, - invalidate_local_hash_match_cache, - search_local_matches_by_hash, - search_local_matches, - ) - from .core.workflow_analyzer import analyze_workflow_models - from .core.path_templates import infer_download_path_templates from .core.metadata_audit import audit_metadata_sizes from .core.metadata_builder import ( build_missing_local_metadata, get_metadata_build_capabilities, ) - from .core.scanner import get_model_files, invalidate_model_files_cache, find_local_file_path - from .core.path_utils import ( - extract_safetensors_header_sha256, - get_filename_from_path, - is_path_in_configured_model_roots, - read_json_safe, - write_json_atomic, - ) from .core.network_utils import ( UnsafeUrlError, host_matches_domain, request_public_url, validate_public_http_url, ) - from .core.type_utils import ( - first_non_empty, - to_int, - to_bool, - format_size_bytes, - normalize_sha256, - extract_sha256_from_metadata, - fetch_remote_file_size_cached, - looks_like_model_file, - normalize_category_to_model_type, - get_category_folder_keys, - get_enabled_download_categories, + from .core.path_templates import infer_download_path_templates + from .core.path_utils import ( + extract_safetensors_header_sha256, + get_filename_from_path, + is_path_in_configured_model_roots, + read_json_safe, + write_json_atomic, ) + from .core.resolver import ( + analyze_and_find_matches, + apply_resolution, + get_local_model_hash_metadata, + invalidate_local_hash_match_cache, + search_local_matches, + search_local_matches_by_hash, + ) + from .core.scanner import find_local_file_path, get_model_files, invalidate_model_files_cache from .core.settings import ( TEMPLATE_KEY_ALIASES, - bool_setting as resolver_bool_setting, get_default_root_for_category, - load_settings as load_resolver_settings, resolve_download_subfolder, + ) + from .core.settings import ( + bool_setting as resolver_bool_setting, + ) + from .core.settings import ( + load_settings as load_resolver_settings, + ) + from .core.settings import ( save_settings as save_resolver_settings, ) + from .core.type_utils import ( + extract_sha256_from_metadata, + fetch_remote_file_size_cached, + first_non_empty, + format_size_bytes, + get_category_folder_keys, + get_enabled_download_categories, + looks_like_model_file, + normalize_category_to_model_type, + normalize_sha256, + to_bool, + to_int, + ) + from .core.workflow_analyzer import analyze_workflow_models except ImportError as e: self.logger.error(f"Model Resolver: Could not import core modules: {e}") return False # Import download modules try: + from .core.aria2_installer import Aria2InstallError, install_aria2_engine from .core.downloader import ( - start_background_download, - get_progress, - get_all_progress, cancel_download, - pause_download, - resume_download, + get_all_progress, get_aria2_status, - start_aria2_daemon, - stop_aria2_daemon, get_download_directory, get_metadata_sidecar_path, + get_progress, get_safe_metadata_sidecar_path, - normalize_download_category, - write_lora_manager_metadata, - sanitize_download_filename, is_allowed_model_download_filename, + normalize_download_category, + pause_download, + resume_download, + sanitize_download_filename, + start_aria2_daemon, + start_background_download, + stop_aria2_daemon, + write_lora_manager_metadata, ) - from .core.aria2_installer import Aria2InstallError, install_aria2_engine - from .core.sources.popular import ( - get_popular_model_url, - search_popular_models, - reload_databases as reload_popular_databases, - ) - from .core.sources.model_list import ( - search_model_list, - search_model_list_multiple, - get_model_list_update_status, - update_model_list_from_remote, - reload_model_list, - ) - from .core.sources.huggingface import ( - search_huggingface_for_file, - parse_huggingface_url, - get_huggingface_download_url, - get_author_fallback_index_status, - refresh_author_fallback_index, - check_huggingface_token, - check_brave_search_api_key, - clear_search_cache as clear_huggingface_search_cache, - ) - from .core.sources.civitai import ( - search_civitai_for_file, - search_civitai, - parse_civitai_url, - get_civitai_download_url, - get_civitai_model_details, - resolve_urn, - check_civitai_api_key, - check_civitai_session_token, - clear_search_cache as clear_civitai_search_cache, - ) + from .core.sources import clear_all_search_caches from .core.sources.civarchive import ( CivArchiveSearchError, + get_civarchive_model_details, is_civarchive_available, - search_civarchive_for_file, parse_civarchive_url, resolve_civarchive_by_hash, resolve_civarchive_model_version, - get_civarchive_model_details, + search_civarchive_for_file, + ) + from .core.sources.civarchive import ( clear_search_cache as clear_civarchive_search_cache, ) + from .core.sources.civitai import ( + check_civitai_api_key, + check_civitai_session_token, + get_civitai_download_url, + get_civitai_model_details, + parse_civitai_url, + resolve_urn, + search_civitai, + search_civitai_for_file, + ) + from .core.sources.civitai import ( + clear_search_cache as clear_civitai_search_cache, + ) + from .core.sources.huggingface import ( + check_brave_search_api_key, + check_huggingface_token, + get_author_fallback_index_status, + get_huggingface_download_url, + parse_huggingface_url, + refresh_author_fallback_index, + search_huggingface_for_file, + ) + from .core.sources.huggingface import ( + clear_search_cache as clear_huggingface_search_cache, + ) + from .core.sources.lora_manager_archive import ( + clear_search_cache as clear_lora_manager_archive_search_cache, + ) from .core.sources.lora_manager_archive import ( is_lora_manager_archive_available, search_lora_manager_archive_for_file, - clear_search_cache as clear_lora_manager_archive_search_cache, ) - from .core.sources import clear_all_search_caches + from .core.sources.model_list import ( + get_model_list_update_status, + reload_model_list, + search_model_list, + search_model_list_multiple, + update_model_list_from_remote, + ) + from .core.sources.popular import ( + get_popular_model_url, + search_popular_models, + ) + from .core.sources.popular import ( + reload_databases as reload_popular_databases, + ) download_available = True except ImportError as e: @@ -1139,6 +1156,7 @@ class ModelResolverExtension: def calculate_sha256_with_progress(normalized_path, progress_id=""): import os as _os + from .core.path_utils import calculate_file_sha256 as _calculate_file_sha256_core total_bytes = max(0, _os.path.getsize(normalized_path)) @@ -1618,8 +1636,8 @@ class ModelResolverExtension: def build_loaded_models_response(): from .core.workflow_analyzer import ( - analyze_workflow_models, URN_TYPE_MAP, + analyze_workflow_models, ) update_loaded_progress( @@ -3452,13 +3470,21 @@ class ModelResolverExtension: if download_available: from .core.path_utils import ( - get_path_abs as get_local_path_abs, - get_path_key as get_local_path_key, - get_path_identity as get_local_path_identity, - get_comfy_root_path, - is_path_within as is_local_path_within, - prefer_local_base_directory, dedupe_local_base_directories, + get_comfy_root_path, + prefer_local_base_directory, + ) + from .core.path_utils import ( + get_path_abs as get_local_path_abs, + ) + from .core.path_utils import ( + get_path_identity as get_local_path_identity, + ) + from .core.path_utils import ( + get_path_key as get_local_path_key, + ) + from .core.path_utils import ( + is_path_within as is_local_path_within, ) # cleanup_search_progress, is_search_progress_cancelled, mark_search_progress_cancelled, and update_search_progress removed @@ -4084,10 +4110,7 @@ class ModelResolverExtension: if model_list_result: confidence = model_list_result.get("confidence", 0) - if is_urn and confidence >= 70: - source_results["model_list"] = model_list_result - source_found = True - elif not is_urn: + if (is_urn and confidence >= 70) or not is_urn: source_results["model_list"] = model_list_result source_found = True @@ -4654,7 +4677,7 @@ class ModelResolverExtension: if not filename: # Extract filename from URL - from urllib.parse import urlparse, unquote + from urllib.parse import unquote, urlparse parsed = urlparse(url) filename = unquote(parsed.path.split("/")[-1]) @@ -4933,6 +4956,7 @@ class ModelResolverExtension: async def get_root_directories(request): """Get configured ComfyUI root directories for path settings.""" import os + import folder_paths known_categories = set(folder_paths.folder_names_and_paths.keys()) @@ -5050,6 +5074,7 @@ class ModelResolverExtension: async def get_subfolders(request): """Get known subfolders for a category using ComfyUI folder_paths.""" import os + import folder_paths raw_category = (request.match_info.get("category") or "").strip() @@ -5278,7 +5303,7 @@ class ModelResolverExtension: lines.append( f"===== {name} ({stat.st_size} bytes, modified {modified_at}) =====" ) - with open(path, "r", encoding="utf-8", errors="replace") as log_file: + with open(path, encoding="utf-8", errors="replace") as log_file: lines.append(log_file.read().rstrip()) except OSError as exc: lines.append(f"===== {name} =====") diff --git a/core/__init__.py b/core/__init__.py index e181279..e8f69a8 100644 --- a/core/__init__.py +++ b/core/__init__.py @@ -11,14 +11,13 @@ Modules: - sources: Search integrations (HuggingFace, CivitAI, popular models) """ +from .matcher import find_matches from .resolver import analyze_and_find_matches, apply_resolution from .scanner import get_model_files -from .matcher import find_matches - __all__ = [ "analyze_and_find_matches", "apply_resolution", - "get_model_files", "find_matches", + "get_model_files", ] diff --git a/core/aria2_installer.py b/core/aria2_installer.py index 9eb493a..5d1b7b9 100644 --- a/core/aria2_installer.py +++ b/core/aria2_installer.py @@ -15,7 +15,6 @@ import requests from .path_utils import is_path_within - ARIA2_RELEASE_API_URL = "https://api.github.com/repos/aria2/aria2/releases/latest" REQUEST_HEADERS = {"User-Agent": "ComfyUI-Model-Resolver"} CUSTOM_NODE_ROOT = Path(__file__).resolve().parents[1] diff --git a/core/catalog_manager.py b/core/catalog_manager.py index 87e3b2b..b27e5fb 100644 --- a/core/catalog_manager.py +++ b/core/catalog_manager.py @@ -1,8 +1,10 @@ """Unified manager for metadata catalog files (JSON data and metadata sidecars).""" from typing import Any, Dict + from .path_utils import read_json_safe, save_catalog_with_backup + class CatalogManager: """Encapsulates reading, writing, status auditing, and backup logic for catalog files.""" diff --git a/core/downloader.py b/core/downloader.py index 6850f3a..09fabc0 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -5,8 +5,6 @@ Handles downloading models from various sources with progress tracking. """ import os -import json -import hashlib import re import secrets import shutil @@ -14,25 +12,44 @@ import socket import subprocess import threading import time -import requests -from typing import Dict, Any, Optional, Callable, List -from pathlib import Path from collections import deque -from urllib.parse import parse_qsl, urlencode, urljoin, urlparse, urlunparse +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional +from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse + +import requests from .log_system import create_module_logger + 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, get_filename_from_path from .network_utils import ( host_matches_domain, request_public_url, validate_public_http_url, ) -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, get_category_folder_keys, MODEL_EXTENSIONS - +from .path_utils import ( + calculate_file_sha256, + get_comfy_root_path, + get_filename_from_path, + get_path_identity, + is_path_within, + read_json_safe, + write_json_atomic, +) +from .resolver import invalidate_local_hash_match_cache, normalize_sha256 +from .scanner import invalidate_model_files_cache +from .type_utils import ( + DEFAULT_BROWSER_USER_AGENT, + MODEL_EXTENSIONS, + as_dict, + as_list, + extract_response_file_size, + first_non_empty, + get_category_folder_keys, + normalize_category_to_model_type, +) +from .type_utils import format_size_bytes as format_bytes try: import folder_paths @@ -72,7 +89,12 @@ HF_XET_ARIA2_AUTH_HOSTS = { "cas-bridge-direct.xethub-eu.hf.co", } -from .settings import CATEGORY_MAP, load_settings, normalize_download_backend, normalize_download_category, normalize_relative_subfolder +from .settings import ( + load_settings, + normalize_download_backend, + normalize_download_category, + normalize_relative_subfolder, +) SENSITIVE_METADATA_KEYS = { "authorization", @@ -1610,7 +1632,7 @@ def _download_huggingface_xet( size = os.path.getsize(partial_path) if expected_size and size != expected_size: - raise IOError( + raise OSError( f"Downloaded size mismatch: expected {expected_size}, received {size}" ) os.replace(partial_path, dest_path) diff --git a/core/log_system/__init__.py b/core/log_system/__init__.py index b968b53..eeac67a 100644 --- a/core/log_system/__init__.py +++ b/core/log_system/__init__.py @@ -4,11 +4,11 @@ AzLogs - Central logging system Provides colored console logging, file rotation, and per-module log levels. """ -from .logger import logger, LogLevel from .log_funcs import create_module_logger +from .logger import LogLevel, logger __all__ = [ - "logger", "LogLevel", "create_module_logger", + "logger", ] diff --git a/core/log_system/log_funcs.py b/core/log_system/log_funcs.py index af7c706..1dc63aa 100644 --- a/core/log_system/log_funcs.py +++ b/core/log_system/log_funcs.py @@ -4,10 +4,9 @@ @version: 2.0.0 @description: Logging Initializer """ -# ruff: noqa: T201 +import logging import os import traceback -import logging _logger = logging.getLogger(__name__) _initialized = False @@ -15,8 +14,8 @@ _project_root = None _default_module_name = __name__ try: - from .logger import logger, LogLevel, debug, info, warn, error, exception, fatal - from .config import LOG_LEVEL, LOG_MODULE_NAME, USE_COLORS, PROJECT_ALIASES + from .config import LOG_LEVEL, LOG_MODULE_NAME, PROJECT_ALIASES, USE_COLORS + from .logger import LogLevel, debug, error, exception, fatal, info, logger, warn def _find_project_root(start_path): current = os.path.dirname(os.path.abspath(start_path)) diff --git a/core/log_system/logger.py b/core/log_system/logger.py index 0e97d94..ef84489 100644 --- a/core/log_system/logger.py +++ b/core/log_system/logger.py @@ -11,17 +11,16 @@ Features: - Log file rotation - Configuration via environment variables """ -# ruff: noqa: T201 -import os -import sys import json -import re import logging -import time +import os +import re +import sys import threading +import time +import traceback from enum import IntEnum from logging.handlers import RotatingFileHandler -import traceback # Log levels @@ -546,9 +545,9 @@ class AzLogsLogger: message = " ".join(str(arg) for arg in args) # Extract known kwargs - exc_info = kwargs.get("exc_info", None) + exc_info = kwargs.get("exc_info") stacklevel = kwargs.get("stacklevel", 1) - extra = kwargs.get("extra", None) + extra = kwargs.get("extra") # Map LogLevel to logging level log_level = LEVEL_MAP.get(level, logging.INFO) diff --git a/core/matcher.py b/core/matcher.py index 3f72e83..291991d 100644 --- a/core/matcher.py +++ b/core/matcher.py @@ -4,13 +4,14 @@ Fuzzy Matcher Module Implements fuzzy string matching to find similar model names. """ +import heapq import os import re -import heapq -from typing import List, Dict, Tuple, Optional, Any from difflib import SequenceMatcher +from typing import Any, Dict, List, Optional, Tuple from .log_system import create_module_logger + log = create_module_logger(__name__) from .path_utils import get_filename_from_path @@ -290,6 +291,7 @@ def base_model_score(candidate: str, preferred: Optional[str]) -> float: # ==================== CENTRALIZED MATCHING & CONFIDENCE HELPERS ==================== from .type_utils import MODEL_EXTENSIONS, PRECISION_FORMAT_SUFFIXES, get_version_sort_key + MODEL_FILE_EXTENSIONS = MODEL_EXTENSIONS @@ -350,11 +352,11 @@ def calculate_filename_confidence(target_filename: str, candidate_filename: str) best_similarity = calculate_similarity_with_normalization( target_filename, candidate_filename ) - + # Cap similarity at 0.999 for non-exact matches to prevent false 100% scores if best_similarity >= 0.999: best_similarity = 0.999 - + return round(best_similarity * 100, 1) @@ -418,26 +420,26 @@ def build_filename_search_queries(filename: str) -> List[str]: basename = get_filename_from_path(filename).strip() if not basename: return [] - + stem = os.path.splitext(basename)[0].strip() stem_lower = stem.lower() - + # Pattern to match common precision/format suffixes pattern = r"[-_]?(" + "|".join(PRECISION_FORMAT_SUFFIXES) + r")" - + # 1. Strip suffix followed by any other characters (like clean_filename_for_search) cleaned = re.sub(pattern + r".*$", "", stem, flags=re.IGNORECASE).strip(" -_") cleaned_lower = cleaned.lower() - + # 2. Strip suffix at the end of the stem (like simplified in search queries) simplified = re.sub(pattern + r"$", "", stem_lower, flags=re.IGNORECASE).strip(" -_") - + queries = [] for query in [basename, stem, stem_lower, cleaned, cleaned_lower, simplified]: query_val = (query or "").strip(" -_") if query_val and query_val not in queries: queries.append(query_val) - + return queries @@ -454,8 +456,7 @@ def match_model_by_title_generic( log_prefix: str = "ModelResolver", ) -> Optional[Dict[str, Any]]: """Generic helper to resolve extensionless workflow values by model title across CivitAI and CivArchive.""" - from typing import Callable, Any - + title_confidence = calculate_model_title_confidence(title_query, model_name) if title_confidence < MODEL_TITLE_MATCH_THRESHOLD: log.debug( @@ -472,7 +473,7 @@ def match_model_by_title_generic( for version in versions: hydrated_version = hydrate_version_fn(version) if hydrate_version_fn else version - + # Check base model context if base_model_context: base_model = get_base_model_fn(hydrated_version) @@ -492,7 +493,7 @@ def match_model_by_title_generic( result["confidence"] = title_confidence result["title_confidence"] = title_confidence - + log.info( f"{log_prefix} model-title match: query={title_query}, model_id={model_id}, model_name={model_name}, version_id={result.get('version_id')}, filename={result.get('filename')}, confidence={title_confidence}, base={result.get('base_model')}" ) diff --git a/core/network_utils.py b/core/network_utils.py index 9fa68c3..06b7df9 100644 --- a/core/network_utils.py +++ b/core/network_utils.py @@ -9,7 +9,6 @@ from urllib.parse import urljoin, urlparse import requests - REDIRECT_STATUS_CODES = {301, 302, 303, 307, 308} SENSITIVE_REDIRECT_HEADERS = {"authorization", "cookie", "proxy-authorization"} @@ -212,6 +211,7 @@ def request_source_response( ) -> Optional[requests.Response]: """Perform a GET/POST request with retry logic for rate limit status (HTTP 429).""" import time + from .log_system import create_module_logger log = create_module_logger(__name__) diff --git a/core/path_templates.py b/core/path_templates.py index 09902f7..f1bc1ad 100644 --- a/core/path_templates.py +++ b/core/path_templates.py @@ -12,7 +12,6 @@ from .settings import ( sanitize_folder_name, ) - TEMPLATE_CATEGORIES = ( "loras", "checkpoints", diff --git a/core/path_utils.py b/core/path_utils.py index c0cb9f6..b14776d 100644 --- a/core/path_utils.py +++ b/core/path_utils.py @@ -5,12 +5,11 @@ Provides helper functions for path normalization, symlink resolution, and directory matching to be shared across all modules. """ -import os -import json import hashlib +import json +import os import re -from typing import Any, Tuple, Optional, Callable, Iterable, List, Dict - +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple _log = None @@ -226,13 +225,13 @@ def write_json_atomic( separators: Optional[Tuple[str, str]] = None, ) -> None: """Safely write data to a JSON file using a temporary file and atomic replace.""" - import tempfile import json - + import tempfile + dir_name = os.path.dirname(file_path) if dir_name: os.makedirs(dir_name, exist_ok=True) - + fd, tmp_path = tempfile.mkstemp( prefix=f".{os.path.basename(file_path)}.", suffix=".tmp", @@ -550,7 +549,7 @@ def _metadata_get(metadata: Dict[str, Any], *keys: str) -> Any: if key in metadata: return metadata.get(key) - lowered = {str(key).lower(): key for key in metadata.keys()} + lowered = {str(key).lower(): key for key in metadata} for key in keys: original_key = lowered.get(str(key).lower()) if original_key is not None: @@ -1001,14 +1000,14 @@ def save_catalog_with_backup( dir_name = os.path.dirname(data_path) if dir_name: os.makedirs(dir_name, exist_ok=True) - + # Backup existing data file if os.path.exists(data_path): shutil.copy2(data_path, f"{data_path}.bak") - + # Write data file atomically write_json_atomic(data_path, data, indent=indent) - + # Write metadata file atomically write_json_atomic(meta_path, meta, indent=indent) @@ -1018,7 +1017,7 @@ def read_json_safe(file_path: str, default: Any = None) -> Any: if not file_path or not os.path.exists(file_path): return default try: - with open(file_path, "r", encoding="utf-8") as f: + with open(file_path, encoding="utf-8") as f: return json.load(f) except Exception as e: _get_log().warning(f"Error reading JSON from {os.path.basename(file_path)}: {e}") @@ -1086,7 +1085,7 @@ def split_path_segments(path_value: Any, filter_dots: bool = True) -> list[str]: text = str(path_value or "").replace("\\", "/").strip() if not text: return [] - + parts = [] for raw_part in text.split("/"): part = raw_part.strip() @@ -1095,6 +1094,6 @@ def split_path_segments(path_value: Any, filter_dots: bool = True) -> list[str]: if filter_dots and part in {".", ".."}: continue parts.append(part) - + return parts diff --git a/core/progress.py b/core/progress.py index f14b522..ae814e6 100644 --- a/core/progress.py +++ b/core/progress.py @@ -4,8 +4,10 @@ Progress Reporting Module Unified utility for reporting progress of asynchronous search and download operations. """ -from typing import Dict, Any, Optional, Callable +from typing import Any, Callable, Dict, Optional + from .log_system import create_module_logger + log = create_module_logger(__name__) diff --git a/core/resolver.py b/core/resolver.py index 014d409..e0f2024 100644 --- a/core/resolver.py +++ b/core/resolver.py @@ -4,17 +4,29 @@ Core Resolver Module Integrates all components to provide high-level API for model linking. """ +import json import os import re -import json import threading -from typing import Dict, Any, List, Optional, Tuple, Callable +from typing import Any, Callable, Dict, List, Optional, Tuple from urllib.parse import unquote from .log_system import create_module_logger + log = create_module_logger(__name__) +from .matcher import find_matches, strip_known_model_extension from .scanner import get_model_files +from .sources.huggingface import parse_huggingface_url as parse_hf_url +from .type_utils import MODEL_EXTENSIONS as _MODEL_EXTENSIONS +from .type_utils import ( + as_dict, + as_list, + extract_sha256_from_metadata, + normalize_sha256, + prepare_remote_size_probe_url, + unique_ordered_strings, +) from .workflow_analyzer import ( NESTED_MODEL_KEYS, NODE_TYPE_TO_CATEGORY_HINTS, @@ -23,11 +35,7 @@ from .workflow_analyzer import ( identify_missing_models, should_scan_as_model_reference, ) -from .matcher import find_matches, strip_known_model_extension -from .type_utils import as_dict, as_list, MODEL_EXTENSIONS as _MODEL_EXTENSIONS, unique_ordered_strings, extract_sha256_from_metadata, normalize_sha256, prepare_remote_size_probe_url from .workflow_updater import update_workflow_nodes -from .sources.civitai import resolve_urn -from .sources.huggingface import parse_huggingface_url as parse_hf_url # Regex patterns for URL extraction (matches HuggingFace and CivitAI URLs) URL_PATTERN = re.compile(r'(https?://(?:huggingface\.co|civitai\.com)[^\s"\'<>\)\\]+)') @@ -37,18 +45,12 @@ MODEL_EXTENSIONS = tuple(_MODEL_EXTENSIONS) from .path_utils import ( - get_path_identity, - get_path_abs, - get_path_key, - is_path_within, - prefer_local_base_directory, - dedupe_local_base_directories, - get_filename_from_path, - read_json_safe, find_metadata_sidecar_path, + get_filename_from_path, + get_path_identity, + read_json_safe, ) - # Imported from .matcher _LOCAL_HASH_MATCH_CACHE_LOCK = threading.Lock() @@ -1305,7 +1307,7 @@ def apply_resolution( } # If resolved_model provided, extract path if needed - if "resolved_model" in resolution and resolution["resolved_model"]: + if resolution.get("resolved_model"): resolved_model = resolution["resolved_model"] if "path" in resolved_model and not mapping.get("resolved_path"): mapping["resolved_path"] = resolved_model["path"] diff --git a/core/scanner.py b/core/scanner.py index 23b07a0..d69b78a 100644 --- a/core/scanner.py +++ b/core/scanner.py @@ -6,13 +6,15 @@ Scans configured model directories and finds available model files. import os import time -from typing import List, Dict, Tuple, Optional +from typing import Dict, List, Optional, Tuple from .log_system import create_module_logger + log = create_module_logger(__name__) -from .path_utils import get_path_identity, get_filename_from_path +from .path_utils import get_filename_from_path, get_path_identity + # Import folder_paths lazily - it may not be available until ComfyUI is initialized try: import folder_paths diff --git a/core/settings.py b/core/settings.py index e0429b1..a020f43 100644 --- a/core/settings.py +++ b/core/settings.py @@ -6,12 +6,12 @@ import json import re from pathlib import Path from typing import Any, Dict, Iterable, Mapping, Optional -from .path_utils import write_json_atomic, read_json_safe +from .path_utils import read_json_safe, write_json_atomic SETTINGS_FILE = Path(__file__).resolve().parents[1] / "model_resolver_settings.json" -from .type_utils import to_bool, CATEGORY_MAP, normalize_download_category +from .type_utils import normalize_download_category, to_bool DOWNLOAD_PATH_MODES = {"suggested", "template", "manual"} DOWNLOAD_BACKENDS = {"python", "aria2"} diff --git a/core/sources/__init__.py b/core/sources/__init__.py index e542672..795ec3b 100644 --- a/core/sources/__init__.py +++ b/core/sources/__init__.py @@ -4,37 +4,45 @@ Model Sources Module Provides search functionality for finding models from various sources. """ -from ..log_system import create_module_logger # noqa: F401 - -from .popular import search_popular_models, get_popular_model_url -from .model_list import search_model_list, search_model_list_multiple -from .huggingface import ( - search_huggingface, - search_huggingface_for_file, - get_huggingface_download_url, - clear_search_cache as clear_huggingface_search_cache, -) -from .civitai import ( - search_civitai, - search_civitai_for_file, - get_civitai_download_url, - clear_search_cache as clear_civitai_search_cache, -) +from ..log_system import create_module_logger from .civarchive import ( clear_search_cache as clear_civarchive_search_cache, +) +from .civarchive import ( is_civarchive_available, resolve_civarchive_by_hash, resolve_civarchive_model_version, search_civarchive, search_civarchive_for_file, ) +from .civitai import ( + clear_search_cache as clear_civitai_search_cache, +) +from .civitai import ( + get_civitai_download_url, + search_civitai, + search_civitai_for_file, +) +from .huggingface import ( + clear_search_cache as clear_huggingface_search_cache, +) +from .huggingface import ( + get_huggingface_download_url, + search_huggingface, + search_huggingface_for_file, +) from .lora_manager_archive import ( clear_search_cache as clear_lora_manager_archive_search_cache, +) +from .lora_manager_archive import ( get_lora_manager_archive_db_path, is_lora_manager_archive_available, search_lora_manager_archive, search_lora_manager_archive_for_file, ) +from .model_list import search_model_list, search_model_list_multiple +from .popular import get_popular_model_url, search_popular_models + def clear_all_search_caches() -> None: """Clear search caches for all external sources.""" @@ -44,26 +52,26 @@ def clear_all_search_caches() -> None: clear_lora_manager_archive_search_cache() __all__ = [ - "search_popular_models", - "get_popular_model_url", - "search_model_list", - "search_model_list_multiple", - "search_huggingface", - "search_huggingface_for_file", - "get_huggingface_download_url", - "search_civitai", - "search_civitai_for_file", - "get_civitai_download_url", + "clear_all_search_caches", "clear_civarchive_search_cache", + "clear_lora_manager_archive_search_cache", + "get_civitai_download_url", + "get_huggingface_download_url", + "get_lora_manager_archive_db_path", + "get_popular_model_url", "is_civarchive_available", + "is_lora_manager_archive_available", "resolve_civarchive_by_hash", "resolve_civarchive_model_version", "search_civarchive", "search_civarchive_for_file", - "clear_lora_manager_archive_search_cache", - "get_lora_manager_archive_db_path", - "is_lora_manager_archive_available", + "search_civitai", + "search_civitai_for_file", + "search_huggingface", + "search_huggingface_for_file", "search_lora_manager_archive", "search_lora_manager_archive_for_file", - "clear_all_search_caches", + "search_model_list", + "search_model_list_multiple", + "search_popular_models", ] diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index da79d36..35cc63e 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -6,49 +6,47 @@ Search and resolve models from CivArchive. import html import json -import os import re import time from typing import Any, Callable, Dict, List, Optional, Tuple -from urllib.parse import parse_qs, unquote, urljoin, urlparse +from urllib.parse import urljoin, urlparse -import requests - -from ..network_utils import host_matches_domain, request_source_response +from ..log_system import create_module_logger from ..matcher import ( + base_model_matches as _base_model_matches, +) +from ..matcher import ( + build_filename_search_queries, calculate_archived_model_confidence, calculate_model_title_confidence, - has_known_model_extension as _has_known_model_extension, - normalize_base_model as _normalize_base_model, - base_model_matches as _base_model_matches, - base_model_score as _base_model_score, - calculate_candidate_rank, - build_filename_search_queries, - MODEL_TITLE_MATCH_THRESHOLD, should_update_best_match, ) +from ..matcher import ( + has_known_model_extension as _has_known_model_extension, +) +from ..matcher import ( + normalize_base_model as _normalize_base_model, +) +from ..network_utils import host_matches_domain, request_source_response from ..path_utils import get_filename_from_path +from ..progress import get_progress_reporter from ..type_utils import ( - to_int, CIVARCHIVE_API_TYPE_MAP, - select_primary_model_file, - parse_size_to_bytes, - get_version_sort_key, DEFAULT_BROWSER_USER_AGENT, - parse_civitai_model_path, - looks_like_model_file, - normalize_model_image, - fetch_remote_file_size_cached, - clear_remote_size_cache, - extract_trained_words, - extract_file_size, - normalize_model_file_info, build_search_result, + clear_remote_size_cache, + extract_file_size, + extract_trained_words, + fetch_remote_file_size_cached, + looks_like_model_file, + normalize_model_file_info, + normalize_model_image, normalize_sha256, prepare_remote_size_probe_url, + select_primary_model_file, + to_int, ) -from ..progress import report_progress, get_progress_reporter -from ..log_system import create_module_logger + log = create_module_logger(__name__) diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 4bb57fc..b7c9975 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -4,40 +4,31 @@ CivitAI Source Module Search and download models from CivitAI. """ +import json import os import re -import json -import hashlib -import requests -from ..network_utils import request_source_response, request_source_json -from typing import Dict, Any, Optional, List, Callable -from urllib.parse import urlparse, parse_qs, quote, urlencode +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 + +from ..log_system import create_module_logger from ..matcher import ( - calculate_filename_confidence, - calculate_model_title_confidence, - has_known_model_extension as _has_known_model_extension, - normalize_base_model as _normalize_base_model, base_model_matches as _base_model_matches, +) +from ..matcher import ( base_model_score as _base_model_score, ) -from ..type_utils import ( - first_non_empty, - CIVITAI_API_TYPE_MAP, - select_primary_model_file, - parse_size_to_bytes, - get_version_sort_key, - check_credential_http, - DEFAULT_BROWSER_USER_AGENT, - parse_civitai_model_path, - normalize_model_image, - as_dict, - as_list, - extract_trained_words, - build_search_result, - normalize_model_file_info, +from ..matcher import ( + calculate_filename_confidence, ) -from ..progress import report_progress, get_progress_reporter +from ..matcher import ( + has_known_model_extension as _has_known_model_extension, +) +from ..matcher import ( + normalize_base_model as _normalize_base_model, +) +from ..network_utils import request_source_json, request_source_response from ..path_utils import ( calculate_file_sha256, extract_safetensors_header_metadata, @@ -46,7 +37,22 @@ from ..path_utils import ( infer_safetensors_base_model, read_json_safe, ) -from ..log_system import create_module_logger +from ..progress import get_progress_reporter +from ..type_utils import ( + CIVITAI_API_TYPE_MAP, + DEFAULT_BROWSER_USER_AGENT, + as_dict, + as_list, + build_search_result, + check_credential_http, + extract_trained_words, + first_non_empty, + normalize_model_file_info, + normalize_model_image, + parse_size_to_bytes, + select_primary_model_file, +) + log = create_module_logger(__name__) diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index b376175..df84b7e 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -4,23 +4,28 @@ HuggingFace Source Module Search and download models from HuggingFace Hub. """ -import json import os import re import threading import time -import requests -from typing import Dict, Any, Optional, List, Callable -from urllib.parse import urlparse, quote +from typing import Any, Callable, Dict, List, Optional +from urllib.parse import quote, urlparse -from ..progress import report_progress, get_progress_reporter from ..log_system import create_module_logger +from ..progress import get_progress_reporter + log = create_module_logger(__name__) -from ..path_utils import write_json_atomic, METADATA_DIR, read_json_safe, get_filename_from_path -from ..network_utils import host_matches_domain, request_source_response, request_source_json -from ..type_utils import check_credential_http, parse_size_header as _parse_size_header, fetch_remote_file_size, fetch_remote_file_size_cached, clear_remote_size_cache, extract_file_size, prepare_remote_size_probe_url -from ..matcher import build_filename_search_queries, clean_filename_for_search +from ..matcher import build_filename_search_queries +from ..network_utils import host_matches_domain, request_source_json, request_source_response +from ..path_utils import METADATA_DIR, get_filename_from_path, read_json_safe, write_json_atomic +from ..type_utils import ( + check_credential_http, + clear_remote_size_cache, + extract_file_size, + fetch_remote_file_size_cached, + prepare_remote_size_probe_url, +) HF_API_URL = "https://huggingface.co/api" HF_AUTHOR_FALLBACKS = ["Comfy-Org"] diff --git a/core/sources/lora_manager_archive.py b/core/sources/lora_manager_archive.py index a4b14ee..a540905 100644 --- a/core/sources/lora_manager_archive.py +++ b/core/sources/lora_manager_archive.py @@ -6,28 +6,35 @@ Search archived CivitAI metadata stored by comfyui-lora-manager. import json import os -import sqlite3 import re +import sqlite3 import time from typing import Any, Callable, Dict, List, Optional +from ..log_system import create_module_logger +from ..matcher import ( + base_model_matches as _base_model_matches, +) +from ..matcher import ( + base_model_score as _base_model_score, +) from ..matcher import ( calculate_archived_model_confidence, - normalize_base_model as _normalize_base_model, - base_model_matches as _base_model_matches, - base_model_score as _base_model_score, calculate_candidate_rank, should_update_best_match, ) +from ..matcher import ( + normalize_base_model as _normalize_base_model, +) from ..path_utils import get_filename_from_path +from ..progress import get_progress_reporter from ..type_utils import ( - select_primary_model_file, + build_search_result, get_generic_filename_tokens, normalize_lora_manager_type, - build_search_result, + select_primary_model_file, ) -from ..progress import report_progress, get_progress_reporter -from ..log_system import create_module_logger + log = create_module_logger(__name__) diff --git a/core/sources/model_list.py b/core/sources/model_list.py index ff1f895..c3965ca 100644 --- a/core/sources/model_list.py +++ b/core/sources/model_list.py @@ -4,22 +4,20 @@ Model List Database Module Search the ComfyUI Manager model-list.json database with fuzzy matching. """ -import os import hashlib -import shutil -import tempfile +import os from datetime import datetime, timezone -from typing import Dict, Any, Optional, List -import requests - -from ..matcher import normalize_filename, calculate_similarity +from typing import Any, Dict, List, Optional from ..log_system import create_module_logger +from ..matcher import calculate_similarity, normalize_filename + log = create_module_logger(__name__) -from ..path_utils import METADATA_DIR from ..catalog_manager import CatalogManager +from ..path_utils import METADATA_DIR + MODEL_LIST_FILE = os.path.join(METADATA_DIR, "model-list.json") MODEL_LIST_META_FILE = os.path.join(METADATA_DIR, "model-list.meta.json") catalog_mgr = CatalogManager(MODEL_LIST_FILE, MODEL_LIST_META_FILE, "models") @@ -30,6 +28,7 @@ MODEL_LIST_GITHUB_API_URL = ( "https://api.github.com/repos/Comfy-Org/ComfyUI-Manager/contents/model-list.json?ref=main" ) from ..network_utils import request_source_json + HTTP_TIMEOUT = 30 # Cache for loaded data @@ -69,7 +68,7 @@ def _get_local_model_list_sha() -> str: try: with open(MODEL_LIST_FILE, "rb") as f: content = f.read() - header = f"blob {len(content)}\0".encode("utf-8") + header = f"blob {len(content)}\0".encode() return hashlib.sha1(header + content).hexdigest() except Exception as e: log.warning(f"Error calculating model-list.json SHA: {e}") diff --git a/core/sources/popular.py b/core/sources/popular.py index 34e39f8..5716f56 100644 --- a/core/sources/popular.py +++ b/core/sources/popular.py @@ -5,20 +5,19 @@ Curated list of common models with known download URLs. """ import os -import json import re -import shutil -import tempfile from datetime import datetime, timezone -from typing import Dict, Any, Optional, List +from typing import Any, Dict, List, Optional from ..log_system import create_module_logger + log = create_module_logger(__name__) -from ..path_utils import METADATA_DIR, read_json_safe -from ..network_utils import request_source_json from ..catalog_manager import CatalogManager +from ..network_utils import request_source_json +from ..path_utils import METADATA_DIR, read_json_safe + POPULAR_MODELS_FILE = os.path.join(METADATA_DIR, "popular-models.json") MODEL_ALIASES_FILE = os.path.join(METADATA_DIR, "model-aliases.json") BASE_MODELS_FILE = os.path.join(METADATA_DIR, "base-models.json") @@ -56,7 +55,6 @@ def _load_model_aliases() -> Dict[str, List[str]]: from ..matcher import normalize_base_model as _normalize_base_model - base_models_mgr = CatalogManager(BASE_MODELS_FILE, BASE_MODELS_META_FILE, "base_models") @@ -212,19 +210,19 @@ def _read_base_models_meta() -> Dict[str, Any]: def generate_aliases(name: str) -> List[str]: aliases = {name.lower()} - + # Replace dots, dashes, underscores with spaces, then normalize spaces normalized = re.sub(r"[\._\-]+", " ", name.lower()) normalized = re.sub(r"\s+", " ", normalized).strip() aliases.add(normalized) version_trimmed = re.sub(r"(\d+)(?:[\s\._\-]+0)+(?!\d)", r"\1", normalized) aliases.add(version_trimmed) - + # Remove all spaces/dashes/dots/underscores collapsed = re.sub(r"[\s\._\-]+", "", name.lower()) aliases.add(collapsed) aliases.add(re.sub(r"[\s\._\-]+", "", version_trimmed)) - + # If there are dots/dashes/underscores, add specific clean replacements if "." in name: aliases.add(name.lower().replace(".", "")) @@ -242,7 +240,7 @@ def generate_aliases(name: str) -> List[str]: first_word = words[0] if len(first_word) > 2 and first_word not in {"stable", "flux", "pony", "sdxl"}: aliases.add(first_word) - + return sorted(list({a for a in aliases if len(a) > 1})) @@ -251,45 +249,45 @@ def get_base_models_status(check_remote: bool = False) -> Dict[str, Any]: local_data = _read_base_models_file() base_models = local_data.get("base_models", []) meta = _read_base_models_meta() - + local_count = len(base_models) local_updated_at = meta.get("updated_at") or "" - + status = { "local_count": local_count, "local_updated_at": local_updated_at, "remote_checked_at": meta.get("last_checked_at") or "", "update_available": False, } - + if check_remote: try: # Fetch from CivitAI API using unified helper enums = request_source_json("https://civitai.com/api/v1/enums", timeout=15, log_name="CivitAI Enums") if enums: remote_models = enums.get("BaseModel", []) - + # Check if there are new models by comparing normalized names and aliases all_existing_normalized = set() for m in base_models: all_existing_normalized.add(_normalize_base_model(m.get("name", ""))) for alias in m.get("aliases", []): all_existing_normalized.add(_normalize_base_model(alias)) - + new_models_found = False for remote_name in remote_models: norm_remote = _normalize_base_model(remote_name) if norm_remote and norm_remote not in all_existing_normalized: new_models_found = True break - + status.update({ "remote_checked_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat(), "update_available": new_models_found }) except Exception as e: log.warning(f"Error checking remote base models: {e}") - + return status @@ -301,33 +299,33 @@ def update_base_models_from_remote() -> Dict[str, Any]: remote_names = enums.get("BaseModel", []) if not remote_names or not isinstance(remote_names, list): raise ValueError("CivitAI enums did not return a valid BaseModel list") - + local_data = _read_base_models_file() base_models = local_data.get("base_models", []) - + # 1. Build a set of all normalized aliases currently in base-models.json all_known_normalized = set() for m in base_models: for alias in m.get("aliases", []): all_known_normalized.add(_normalize_base_model(alias)) all_known_normalized.add(_normalize_base_model(m.get("name", ""))) - + updated_models = list(base_models) new_added_count = 0 new_added_names = [] - + for name in remote_names: norm_name = _normalize_base_model(name) if not norm_name: continue - + # If the normalized name is already a known alias or name, skip it if norm_name in all_known_normalized: continue - + # Generate candidate aliases candidates = generate_aliases(name) - + # Filter out aliases that are already registered/known filtered_aliases = [] for alias in candidates: @@ -335,20 +333,20 @@ def update_base_models_from_remote() -> Dict[str, Any]: if norm_alias not in all_known_normalized: filtered_aliases.append(alias) all_known_normalized.add(norm_alias) - + # If we have valid aliases, add the new base model entry if filtered_aliases: if name.lower() not in [a.lower() for a in filtered_aliases]: filtered_aliases.insert(0, name) all_known_normalized.add(norm_name) - + updated_models.append({ "name": name, "aliases": filtered_aliases }) new_added_count += 1 new_added_names.append(name) - + now_str = datetime.now(timezone.utc).replace(microsecond=0).isoformat() meta = { "updated_at": now_str, @@ -359,7 +357,7 @@ def update_base_models_from_remote() -> Dict[str, Any]: } base_models_mgr.save({"base_models": updated_models}, meta, indent=2) reload_databases() - + return { "local_count": len(updated_models), "local_updated_at": now_str, diff --git a/core/type_utils.py b/core/type_utils.py index 05bad61..671d83e 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -262,8 +262,9 @@ CIVARCHIVE_API_TYPE_MAP = { "motion_modules": "MotionModule", } -from typing import Any, Dict, List, Optional, Tuple, Callable import re +from typing import Any, Callable, Dict, List, Optional, Tuple + def normalize_download_category(category: str) -> str: """Return the canonical ComfyUI folder_paths key for a download category.""" @@ -561,7 +562,7 @@ def fetch_remote_file_size( """ from .network_utils import request_public_url request_headers = {**(headers or {}), "Accept-Encoding": "identity"} - + size = None try: response, _final_url, _final_headers = request_public_url( @@ -604,6 +605,7 @@ def looks_like_model_file(url: str, expected_filename: str = "") -> bool: """ import os from urllib.parse import unquote, urlparse + from .network_utils import host_matches_domain from .path_utils import get_filename_from_path @@ -620,7 +622,7 @@ def looks_like_model_file(url: str, expected_filename: str = "") -> bool: if host_matches_domain(host, "huggingface.co") and path.startswith("/spaces/"): return False - + # Civitai/CivArchive api/download url prefix check if ( host_matches_domain(host, "civitai.com") @@ -650,7 +652,7 @@ def normalize_model_image(image_data: Dict[str, Any], default_civitai_url: str = meta = {} url = image_data.get("url") or image_data.get("imageUrl") or image_data.get("src") or "" - + civitai_url = image_data.get("civitaiUrl") or image_data.get("postUrl") or default_civitai_url if not civitai_url and isinstance(image_data.get("id"), (int, str)): civitai_url = f"https://civitai.com/images/{image_data.get('id')}" @@ -693,7 +695,6 @@ def normalize_model_image(image_data: Dict[str, Any], default_civitai_url: str = } -import re SHA256_PATTERN = re.compile(r"^[a-fA-F0-9]{64}$") @@ -836,6 +837,7 @@ def prepare_remote_size_probe_url(url: Any) -> Optional[str]: Converts HuggingFace '/blob/' URLs to '/resolve/' URLs. """ from urllib.parse import urlparse + from .network_utils import host_matches_domain if not isinstance(url, str) or not url.strip(): @@ -992,7 +994,7 @@ def get_enabled_download_categories(folder_names: list[str]) -> list[str]: "text_encoders", "ipadapter", "sams", "ultralytics" ] skip = {"custom_nodes", "configs"} - + categories = [] for cat in [*preferred, *folder_names]: norm = normalize_download_category(cat) @@ -1105,12 +1107,12 @@ def build_search_result( def parse_provider_model_url(url: str, allowed_domains: list[str]) -> Optional[Dict[str, Any]]: """Parse a provider model URL (e.g. CivitAI or mirror site) to extract model/version info.""" - from urllib.parse import urlparse import re - + from urllib.parse import urlparse + if not isinstance(url, str) or not url.strip(): return None - + parsed = urlparse(url) hostname = parsed.hostname if hostname: @@ -1122,7 +1124,7 @@ def parse_provider_model_url(url: str, allowed_domains: list[str]) -> Optional[D break if not matched: return None - + # Standard CivArchive sha256 mirrors pattern sha_match = re.search(r"/sha256/([a-fA-F0-9]{64})", parsed.path) if sha_match: @@ -1133,6 +1135,6 @@ def parse_provider_model_url(url: str, allowed_domains: list[str]) -> Optional[D match = re.search(r"/api/download/models/(\d+)", parsed.path) if match: return {"version_id": int(match.group(1))} - + return parse_civitai_model_path(parsed.path, parsed.query) diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py index 6f11658..ec7bf78 100644 --- a/core/workflow_analyzer.py +++ b/core/workflow_analyzer.py @@ -7,9 +7,10 @@ Extracts model references from workflow JSON and identifies missing models. import os import re import threading -from typing import List, Dict, Any, Optional, Set, Callable +from typing import Any, Callable, Dict, List, Optional, Set from .log_system import create_module_logger + log = create_module_logger(__name__) @@ -21,7 +22,7 @@ except ImportError: log.warning("Model Resolver: folder_paths not available yet - will retry later") -from .type_utils import URN_TYPE_MAP, MODEL_EXTENSIONS, unique_ordered_strings +from .type_utils import MODEL_EXTENSIONS, URN_TYPE_MAP, unique_ordered_strings URN_REGEX = re.compile(r"^urn:air:([^:]+):([^:]+):([^:]+):(\d+)@(\d+)$") @@ -734,7 +735,7 @@ def _build_dynamic_node_widget_category_hints(node_type: str) -> Dict[str, Any]: try: if folder_paths_module is not None and callable(get_filename_list): - setattr(folder_paths_module, "get_filename_list", traced_get_filename_list) + folder_paths_module.get_filename_list = traced_get_filename_list patched_folder_paths = True patch_get_filename_list_global(input_types_func) patch_get_filename_list_global(schema_func) @@ -750,7 +751,7 @@ def _build_dynamic_node_widget_category_hints(node_type: str) -> Dict[str, Any]: for global_scope, global_name, global_value in patched_globals: global_scope[global_name] = global_value if patched_folder_paths: - setattr(folder_paths_module, "get_filename_list", get_filename_list) + folder_paths_module.get_filename_list = get_filename_list if not isinstance(input_types, dict) and schema is None: return empty_hints @@ -1037,7 +1038,7 @@ def get_model_widget_category_hints( if indexed_category_hint: hints_list.append(indexed_category_hint) - + # If the widget category hint is just "diffusion_models", and we have a more specific output hint, # let the output hint take priority. if widget_category_hint == "diffusion_models" and output_widget_category_hint and output_widget_category_hint != "diffusion_models": @@ -1048,7 +1049,7 @@ def get_model_widget_category_hints( hints_list.append(widget_category_hint) if output_widget_category_hint: hints_list.append(output_widget_category_hint) - + return _ordered_unique_categories(hints_list) diff --git a/core/workflow_updater.py b/core/workflow_updater.py index 28540bc..6c6a053 100644 --- a/core/workflow_updater.py +++ b/core/workflow_updater.py @@ -5,12 +5,12 @@ Updates workflow JSON by replacing model paths in nodes. """ import os -from typing import Dict, Any, List, Optional +from typing import Any, Dict, List, Optional from .log_system import create_module_logger -log = create_module_logger(__name__) -from .path_utils import is_path_within, get_filename_from_path +log = create_module_logger(__name__) +from .path_utils import get_filename_from_path, is_path_within def convert_to_relative_path( diff --git a/pyproject.toml b/pyproject.toml index 7ff1b32..cb76adb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,3 +19,49 @@ PublisherId = "azornes" DisplayName = "Comfyui Model Resolver" Icon = "" requires-comfyui = ">=0.26.0" + +# --------------------------------------------------------------------------- +# Ruff – static analysis for Python +# --------------------------------------------------------------------------- +[tool.ruff] +target-version = "py310" +line-length = 120 + +[tool.ruff.lint] +# Rule sets: +# E/W – pycodestyle errors & warnings +# F – pyflakes +# I – isort (import sorting) +# B – flake8-bugbear +# SIM – flake8-simplify +# UP – pyupgrade +# RUF – Ruff-specific rules +select = ["E", "W", "F", "I", "B", "SIM", "UP", "RUF"] + +ignore = [ + "E501", # line-too-long – handled by formatter / not enforced + "E402", # module-import-not-at-top – ComfyUI entry point needs late imports + "W293", # blank-line-with-whitespace – CRLF artifacts on indented blank lines + "B023", # function-uses-loop-variable – intentional in callback patterns + "SIM105", # suppressible-exception – contextlib.suppress not always clearer + "SIM108", # if-else-block-instead-of-if-exp – readability preference + "UP006", # non-pep585-annotation – project targets >=3.10 but uses typing compat + "UP035", # deprecated-import – typing compat kept for clarity + "RUF013", # implicit-optional – existing API signatures + "UP045", # non-pep604-annotation-optional – existing code uses Optional[X] + "UP038", # non-pep604-isinstance – existing isinstance(x, (A, B)) patterns +] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = [ + "F401", # unused-import – re-exports + "E402", # module-import-not-at-top – ComfyUI entry point structure + "F811", # redefined-while-unused – re-exports / conditional imports +] +"tests/*" = [ + "B011", # assert-false – ok in tests + "SIM300", # yoda-conditions – ok in tests +] + +[tool.ruff.lint.isort] +known-first-party = ["core"] diff --git a/tests/test_huggingface.py b/tests/test_huggingface.py index 7f4d37e..392667a 100644 --- a/tests/test_huggingface.py +++ b/tests/test_huggingface.py @@ -8,10 +8,9 @@ from core.sources.huggingface import ( parse_huggingface_url, get_huggingface_download_url, _normalize_huggingface_size_probe_url, - clean_filename_for_search, _write_persistent_author_index, ) -from core.matcher import build_filename_search_queries +from core.matcher import build_filename_search_queries, clean_filename_for_search from core.type_utils import extract_file_size class HuggingFaceSourceTests(unittest.TestCase):