From 65ff2baa300ddd8590416c3ddbb2f0136ad47e89 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Fri, 3 Jul 2026 17:09:35 +0200 Subject: [PATCH] Add local hash match cache and invalidate it on rescan/downloads --- __init__.py | 8 +++ core/downloader.py | 7 ++- core/resolver.py | 154 +++++++++++++++++++++++++++++++-------------- 3 files changed, 120 insertions(+), 49 deletions(-) diff --git a/__init__.py b/__init__.py index 597cb4e..0c0b89a 100644 --- a/__init__.py +++ b/__init__.py @@ -183,6 +183,7 @@ class ModelResolverExtension: analyze_and_find_matches, apply_resolution, get_local_model_hash_metadata, + invalidate_local_hash_match_cache, search_local_matches_by_hash, search_local_matches, ) @@ -351,6 +352,8 @@ class ModelResolverExtension: if isinstance(force_rescan, bool) else str(force_rescan).lower() in {"1", "true", "yes"} ) + if force_rescan: + invalidate_local_hash_match_cache() if workflow_json is None: return web.json_response( @@ -591,6 +594,8 @@ class ModelResolverExtension: filename = data.get("filename", "") category = data.get("category", "") force_rescan = to_bool(data.get("force_rescan"), False) + if force_rescan: + invalidate_local_hash_match_cache() if not filename: return web.json_response( @@ -1018,6 +1023,8 @@ class ModelResolverExtension: or request.query.get("force_rescan") or "" ).lower() in {"1", "true", "yes"} + if force_rescan: + invalidate_local_hash_match_cache() models = get_model_files(force_rescan=force_rescan) return web.json_response(models) @@ -3381,6 +3388,7 @@ class ModelResolverExtension: reload_popular_databases() reload_model_list() invalidate_model_files_cache() + invalidate_local_hash_match_cache() self.search_result_timestamps.clear() self.logger.info("Cleared backend search caches") return web.json_response({"success": True, "cleared": "all"}) diff --git a/core/downloader.py b/core/downloader.py index 37a6229..e180037 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -22,7 +22,8 @@ from urllib.parse import parse_qsl, urlencode, urljoin, urlparse, urlunparse from .log_system import create_module_logger log = create_module_logger(__name__) -from .resolver import normalize_sha256 +from .resolver import invalidate_local_hash_match_cache, normalize_sha256 +from .scanner import invalidate_model_files_cache from .path_utils import is_path_within, get_path_identity, write_json_atomic, read_json_safe, get_comfy_root_path, calculate_file_sha256 as _calculate_file_sha256, get_filename_from_path from .type_utils import as_dict, as_list, first_non_empty, format_size_bytes as format_bytes, DEFAULT_BROWSER_USER_AGENT, extract_response_file_size @@ -1324,6 +1325,8 @@ def download_file_with_aria2( log.info( f"Size: {format_bytes(size)}, Time: {elapsed:.1f}s, Avg speed: {format_bytes(int(avg_speed))}/s" ) + invalidate_model_files_cache() + invalidate_local_hash_match_cache() return result if state == "error": @@ -1599,6 +1602,8 @@ def download_file( log.info( f"Size: {format_bytes(downloaded)}, Time: {elapsed:.1f}s, Avg speed: {format_bytes(int(avg_speed))}/s" ) + invalidate_model_files_cache() + invalidate_local_hash_match_cache() except requests.exceptions.RequestException as e: error_msg = str(e) diff --git a/core/resolver.py b/core/resolver.py index 19b56c2..2f671ab 100644 --- a/core/resolver.py +++ b/core/resolver.py @@ -7,6 +7,7 @@ Integrates all components to provide high-level API for model linking. import os import re import json +import threading from typing import Dict, Any, List, Optional, Tuple, Callable from urllib.parse import unquote @@ -50,6 +51,102 @@ from .path_utils import ( # Imported from .matcher +_LOCAL_HASH_MATCH_CACHE_LOCK = threading.Lock() +_LOCAL_HASH_MATCH_CACHE: Optional[Dict[str, List[Dict[str, Any]]]] = None + + +def invalidate_local_hash_match_cache() -> None: + """Clear the in-memory SHA256 -> local model match index.""" + global _LOCAL_HASH_MATCH_CACHE + with _LOCAL_HASH_MATCH_CACHE_LOCK: + _LOCAL_HASH_MATCH_CACHE = None + + +def _clone_hash_match(match: Dict[str, Any]) -> Dict[str, Any]: + cloned = dict(match) + if isinstance(cloned.get("model"), dict): + cloned["model"] = dict(cloned["model"]) + return cloned + + +def _build_local_hash_match_cache( + available_models: List[Dict[str, Any]], +) -> Dict[str, List[Dict[str, Any]]]: + index: Dict[str, List[Dict[str, Any]]] = {} + seen_entries = set() + + for model in available_models: + model_path = model.get("path", "") + if not model_path: + continue + + metadata_path = find_metadata_sidecar_path(model_path) + if not metadata_path: + continue + + metadata = read_json_safe(metadata_path, None) + if not isinstance(metadata, dict) or not metadata: + log.debug(f"Could not read metadata sidecar for hash match: {metadata_path}") + continue + + metadata_hashes = _extract_model_sha256_from_metadata(metadata, model) + if not metadata_hashes: + continue + + model_filename = model.get("filename") or get_filename_from_path(model_path) + for metadata_hash in metadata_hashes: + normalized_hash = normalize_sha256(metadata_hash) + if not normalized_hash: + continue + + try: + model_identity = get_path_identity(model_path) + except (OSError, ValueError): + model_identity = os.path.normcase(os.path.abspath(model_path)) + entry_key = (normalized_hash, model_identity or model_path) + if entry_key in seen_entries: + continue + seen_entries.add(entry_key) + + model_with_metadata = { + **model, + "sha256": normalized_hash, + "metadata_path": metadata_path, + } + index.setdefault(normalized_hash, []).append( + { + "model": model_with_metadata, + "filename": model_filename, + "similarity": 1.0, + "confidence": 100.0, + "match_type": "hash", + "hash_match": True, + "hash_source": "metadata", + "sha256": normalized_hash, + "metadata_path": metadata_path, + } + ) + + return index + + +def _get_local_hash_match_cache(force_rescan: bool = False) -> Dict[str, List[Dict[str, Any]]]: + global _LOCAL_HASH_MATCH_CACHE + if force_rescan: + invalidate_local_hash_match_cache() + + with _LOCAL_HASH_MATCH_CACHE_LOCK: + if _LOCAL_HASH_MATCH_CACHE is not None: + return _LOCAL_HASH_MATCH_CACHE + + available_models = get_model_files(force_rescan=force_rescan) + index = _build_local_hash_match_cache(available_models) + + with _LOCAL_HASH_MATCH_CACHE_LOCK: + if _LOCAL_HASH_MATCH_CACHE is None: + _LOCAL_HASH_MATCH_CACHE = index + return _LOCAL_HASH_MATCH_CACHE + def get_workflow_url_info_for_filename( workflow_urls: Dict[str, Dict[str, Any]], filename: str @@ -351,57 +448,18 @@ def search_local_matches_by_hash( if not normalized_hash: return [] - available_models = get_model_files(force_rescan=force_rescan) + index = _get_local_hash_match_cache(force_rescan=force_rescan) + matches = [_clone_hash_match(match) for match in index.get(normalized_hash, [])] + if category and category != "unknown": - preferred_models = [ - model for model in available_models if model.get("category") == category - ] - other_models = [ - model for model in available_models if model.get("category") != category - ] - available_models = preferred_models + other_models - - matches: List[Dict[str, Any]] = [] - - for model in available_models: - model_path = model.get("path", "") - if not model_path: - continue - - metadata_path = find_metadata_sidecar_path(model_path) - if not metadata_path: - continue - - metadata = read_json_safe(metadata_path, None) - if not isinstance(metadata, dict) or not metadata: - log.debug(f"Could not read metadata sidecar for hash match: {metadata_path}") - continue - - metadata_hashes = _extract_model_sha256_from_metadata(metadata, model) - if normalized_hash not in metadata_hashes: - continue - - model_with_metadata = { - **model, - "sha256": normalized_hash, - "metadata_path": metadata_path, - } - matches.append( - { - "model": model_with_metadata, - "filename": model.get("filename") or get_filename_from_path(model_path), - "similarity": 1.0, - "confidence": 100.0, - "match_type": "hash", - "hash_match": True, - "hash_source": "metadata", - "sha256": normalized_hash, - "metadata_path": metadata_path, - } + matches.sort( + key=lambda match: 0 + if match.get("model", {}).get("category") == category + else 1 ) - if max_matches > 0 and len(matches) >= max_matches: - break + if max_matches > 0: + return matches[:max_matches] return matches