From 27ce316c67bcf43caca00f099e0e0887d1d85c8d Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Tue, 8 Sep 2026 18:43:51 +0200 Subject: [PATCH] perf: optimize hash cache building with in-memory dir entries and thread pool --- core/path_utils.py | 99 +++++++++++++++++++++++----- core/resolver.py | 61 +++++++++++++---- tests/test_metadata_model_helpers.py | 59 +++++++++++++++++ 3 files changed, 190 insertions(+), 29 deletions(-) diff --git a/core/path_utils.py b/core/path_utils.py index 67724a3..7fc99cb 100644 --- a/core/path_utils.py +++ b/core/path_utils.py @@ -9,7 +9,7 @@ import hashlib import json import os import re -from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple +from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple from .type_utils import normalize_alphanumeric_key @@ -1080,7 +1080,34 @@ def get_comfy_root_path(folder_paths_module: Optional[Any] = None) -> str: return "" -def find_external_metadata_sidecar_path(model_path: str) -> str: +def get_directory_entry_set( + directory: str, + cache: Optional[Dict[str, Set[str]]] = None, +) -> Set[str]: + """ + Return the set of case-normalized filenames in a directory. + + If ``cache`` is provided, it will be checked and updated with the result + keyed by normalized directory path. + """ + if not directory: + return set() + normalized_dir = os.path.normcase(os.path.abspath(directory)) + if cache is not None and normalized_dir in cache: + return cache[normalized_dir] + try: + entries = {os.path.normcase(name) for name in os.listdir(normalized_dir)} + except OSError: + entries = set() + if cache is not None: + cache[normalized_dir] = entries + return entries + + +def find_external_metadata_sidecar_path( + model_path: str, + dir_entries: Optional[Set[str]] = None, +) -> str: """ Find metadata owned by another tool (for example .metadata.json, .civitai.info, or .json). These files are read-only to Model Resolver. @@ -1094,7 +1121,10 @@ def find_external_metadata_sidecar_path(model_path: str) -> str: # Base name without extension base_name = filename.rsplit(".", 1)[0] if "." in filename else filename - has_stem_collision = has_model_sidecar_name_collision(model_path) + has_stem_collision = has_model_sidecar_name_collision( + model_path, + dir_entries=dir_entries, + ) # Name patterns based on civitai.py and resolver.py. exact_names = [ @@ -1113,14 +1143,17 @@ def find_external_metadata_sidecar_path(model_path: str) -> str: exact_names, legacy_names, has_stem_collision=has_stem_collision, + dir_entries=dir_entries, ) for name in possible_names: if name: + if dir_entries is not None and os.path.normcase(name) not in dir_entries: + continue path = os.path.join(directory, name) if os.path.normcase(os.path.abspath(path)) == normalized_model_path: continue - if not os.path.exists(path): + if dir_entries is None and not os.path.exists(path): continue if ( has_stem_collision @@ -1140,10 +1173,17 @@ def get_model_resolver_sidecar_path(model_path: str) -> str: return os.path.join(directory, f"{filename}{MODEL_RESOLVER_METADATA_SUFFIX}") -def find_model_resolver_sidecar_path(model_path: str) -> str: +def find_model_resolver_sidecar_path( + model_path: str, + dir_entries: Optional[Set[str]] = None, +) -> str: """Return the existing Model Resolver sidecar path, if present.""" if not model_path: return "" + filename = get_filename_from_path(model_path) + expected_name = f"{filename}{MODEL_RESOLVER_METADATA_SUFFIX}" + if dir_entries is not None and os.path.normcase(expected_name) not in dir_entries: + return "" metadata_path = get_model_resolver_sidecar_path(model_path) normalized_model_path = os.path.normcase(os.path.abspath(model_path)) if ( @@ -1154,7 +1194,10 @@ def find_model_resolver_sidecar_path(model_path: str) -> str: return "" -def find_metadata_sidecar_path(model_path: str) -> str: +def find_metadata_sidecar_path( + model_path: str, + dir_entries: Optional[Set[str]] = None, +) -> str: """ Return the preferred readable sidecar path. @@ -1162,8 +1205,8 @@ def find_metadata_sidecar_path(model_path: str) -> str: available as a read-only fallback. """ return ( - find_model_resolver_sidecar_path(model_path) - or find_external_metadata_sidecar_path(model_path) + find_model_resolver_sidecar_path(model_path, dir_entries=dir_entries) + or find_external_metadata_sidecar_path(model_path, dir_entries=dir_entries) ) @@ -1242,7 +1285,10 @@ def split_path_segments(path_value: Any, filter_dots: bool = True) -> list[str]: return parts -def has_model_sidecar_name_collision(model_path: str) -> bool: +def has_model_sidecar_name_collision( + model_path: str, + dir_entries: Optional[Set[str]] = None, +) -> bool: """Return whether another model file shares this file's stem.""" if not model_path: return False @@ -1256,12 +1302,20 @@ def has_model_sidecar_name_collision(model_path: str) -> bool: from .type_utils import MODEL_EXTENSIONS normalized_model_path = os.path.normcase(os.path.abspath(model_path)) + normalized_filename = os.path.normcase(filename) for candidate_extension in MODEL_EXTENSIONS: - candidate_path = os.path.join(directory, f"{base_name}{candidate_extension}") - if os.path.normcase(os.path.abspath(candidate_path)) == normalized_model_path: + candidate_filename = f"{base_name}{candidate_extension}" + if os.path.normcase(candidate_filename) == normalized_filename: continue - if os.path.isfile(candidate_path): - return True + if dir_entries is not None: + if os.path.normcase(candidate_filename) in dir_entries: + return True + else: + candidate_path = os.path.join(directory, candidate_filename) + if os.path.normcase(os.path.abspath(candidate_path)) == normalized_model_path: + continue + if os.path.isfile(candidate_path): + return True return False @@ -1329,18 +1383,27 @@ def _ordered_sidecar_names( legacy_names: List[Optional[str]], *, has_stem_collision: Optional[bool] = None, + dir_entries: Optional[Set[str]] = None, ) -> List[str]: """Return sidecar names in the established exact/legacy preference order.""" directory = os.path.dirname(model_path) exact_names = [name for name in exact_names if name] legacy_names = [name for name in legacy_names if name] if has_stem_collision is None: - has_stem_collision = has_model_sidecar_name_collision(model_path) + has_stem_collision = has_model_sidecar_name_collision( + model_path, + dir_entries=dir_entries, + ) - exact_path = os.path.join(directory, exact_names[0]) if exact_names else "" - prefer_exact_names = has_stem_collision or ( - bool(exact_path) and os.path.isfile(exact_path) - ) + exact_exists = False + if exact_names: + if dir_entries is not None: + exact_exists = os.path.normcase(exact_names[0]) in dir_entries + else: + exact_path = os.path.join(directory, exact_names[0]) + exact_exists = bool(exact_path) and os.path.isfile(exact_path) + + prefer_exact_names = has_stem_collision or exact_exists if prefer_exact_names: return exact_names + legacy_names diff --git a/core/resolver.py b/core/resolver.py index ccc6f15..186b52b 100644 --- a/core/resolver.py +++ b/core/resolver.py @@ -8,8 +8,9 @@ import json import os import re import threading +from concurrent.futures import ThreadPoolExecutor from dataclasses import replace -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Set, Tuple from urllib.parse import unquote from .contracts import ( @@ -63,6 +64,7 @@ MODEL_EXTENSIONS = tuple(_MODEL_EXTENSIONS) from .path_utils import ( find_metadata_sidecar_path, + get_directory_entry_set, get_filename_from_path, get_path_identity, get_path_key, @@ -193,36 +195,73 @@ def _build_local_hash_match_cache( seen_entries = set() normalized_models = normalize_models(available_models) + candidates = [ + model + for model in normalized_models + if model.path and _is_local_hash_match_candidate(model) + ] + if not candidates: + return index - for model in normalized_models: - if not _is_local_hash_match_candidate(model): - continue + dir_cache: Dict[str, Set[str]] = {} + for model in candidates: + model_path = model.path + if model_path: + directory = os.path.dirname(model_path) + if directory: + get_directory_entry_set(directory, cache=dir_cache) + def _inspect_model_metadata( + model: ResolvedModel, + ) -> List[Tuple[str, str, ResolvedModel, str, str]]: model_path = model.path if not model_path: - continue + return [] - metadata_path = find_metadata_sidecar_path(model_path) + directory = os.path.dirname(model_path) + dir_entries = dir_cache.get(os.path.normcase(os.path.abspath(directory))) if directory else None + metadata_path = find_metadata_sidecar_path(model_path, dir_entries=dir_entries) if not metadata_path: - continue + return [] metadata = read_merged_model_metadata(model_path, None) if not isinstance(metadata, dict) or not metadata: log.debug(f"Could not read metadata sidecar for hash match: {metadata_path}") - continue + return [] metadata_hashes = _extract_model_sha256_from_metadata(metadata, model) if not metadata_hashes: - continue + return [] model_filename = model.filename or get_filename_from_path(model_path) + model_identity = get_path_identity(model_path) or model_path + results: List[Tuple[str, str, ResolvedModel, str, str]] = [] for metadata_hash in metadata_hashes: normalized_hash = normalize_sha256(metadata_hash) if not normalized_hash: continue + results.append( + ( + normalized_hash, + model_identity, + model, + model_filename, + metadata_path, + ) + ) + return results - model_identity = get_path_identity(model_path) - entry_key = (normalized_hash, model_identity or model_path) + if len(candidates) <= 1: + inspected_results = [_inspect_model_metadata(candidates[0])] + else: + cpu_cores = os.cpu_count() or 2 + max_workers = min(8, max(2, cpu_cores)) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + inspected_results = list(executor.map(_inspect_model_metadata, candidates)) + + for item_list in inspected_results: + for normalized_hash, model_identity, model, model_filename, metadata_path in item_list: + entry_key = (normalized_hash, model_identity) if entry_key in seen_entries: continue seen_entries.add(entry_key) diff --git a/tests/test_metadata_model_helpers.py b/tests/test_metadata_model_helpers.py index 79f6c1b..26987c8 100644 --- a/tests/test_metadata_model_helpers.py +++ b/tests/test_metadata_model_helpers.py @@ -120,6 +120,65 @@ class MetadataModelHelperTests(unittest.TestCase): self.assertEqual("", get_model_path_identity("")) self.assertEqual("", get_model_path_identity(" ")) + def test_build_local_hash_match_cache_parallel_execution(self): + model1 = os.path.normpath("e:/models/checkpoints/model1.safetensors") + model2 = os.path.normpath("e:/models/loras/model2.safetensors") + + def mock_sidecar(path, dir_entries=None): + if "model1" in path: + return path + ".metadata.json" + if "model2" in path: + return path + ".metadata.json" + return "" + + def mock_read(path, default=None): + if "model1" in path: + return {"sha256": "1" * 64} + if "model2" in path: + return {"sha256": "2" * 64} + return default + + def mock_extract(metadata, model): + return [metadata.get("sha256")] if metadata.get("sha256") else [] + + with patch("core.resolver.os.path.isdir", return_value=False), patch( + "core.resolver.find_metadata_sidecar_path", side_effect=mock_sidecar + ), patch( + "core.resolver.read_merged_model_metadata", side_effect=mock_read + ), patch( + "core.resolver._extract_model_sha256_from_metadata", side_effect=mock_extract + ): + index = _build_local_hash_match_cache( + [ + {"path": model1, "filename": "model1.safetensors"}, + {"path": model2, "filename": "model2.safetensors"}, + ] + ) + + self.assertIn("1" * 64, index) + self.assertIn("2" * 64, index) + self.assertEqual(model1, index["1" * 64][0].model.path) + self.assertEqual(model2, index["2" * 64][0].model.path) + + def test_sidecar_lookup_with_dir_entries(self): + from core.path_utils import ( + find_metadata_sidecar_path, + has_model_sidecar_name_collision, + ) + + model_path = "e:/models/test_model.safetensors" + dir_entries = {"test_model.safetensors", "test_model.safetensors.modelresolver.json"} + + with patch("os.path.isfile", return_value=True): + sidecar = find_metadata_sidecar_path(model_path, dir_entries=dir_entries) + self.assertTrue(sidecar.endswith("test_model.safetensors.modelresolver.json")) + + collision_entries = {"test_model.safetensors", "test_model.ckpt"} + self.assertTrue(has_model_sidecar_name_collision(model_path, dir_entries=collision_entries)) + + no_collision_entries = {"test_model.safetensors", "other_model.ckpt"} + self.assertFalse(has_model_sidecar_name_collision(model_path, dir_entries=no_collision_entries)) + if __name__ == "__main__": unittest.main()