From e096f2c786722b702e0ee6500847c2c1b5821c8f Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Fri, 24 Apr 2026 23:32:00 +0200 Subject: [PATCH] Optimize model matching, caching, and CivitAI URLs Performance and correctness improvements across model scanning, matching, workflow analysis, and CivitAI URL handling: - Add short-lived in-memory cache for get_model_files (with force_rescan and invalidate_model_files_cache) to reduce frequent directory scans. - Pass scanned available_models into workflow analysis to avoid rescanning and build category-indexed lists plus an ordered-candidates cache for faster candidate ordering. - Improve analyze_and_find_matches to special-case urn:air: URNs (use filename suffix when present) and skip fuzzy matching when appropriate. - Rewrite find_matches to maintain only the top N results via heapq, cache normalized candidate path/filename values on candidate dicts, and early-exit on max_results<=0 to reduce allocations and CPU. - Add helper functions in CivitAI source to extract/build stable image page URLs (_extract_civitai_image_id, _build_civitai_image_url) and use them when populating image metadata. - Propagate available_models into get_node_model_info and analyze_workflow_models so subroutines reuse the same scanned list. - Guard frontend and backend folder lookups to return early for unknown categories to avoid unnecessary work. These changes aim to reduce IO/CPU during frequent matching operations and produce more accurate/efficient matching for URNs and local models. --- __init__.py | 10 +++++ core/linker.py | 40 ++++++++++++++++--- core/matcher.py | 83 ++++++++++++++++++++++++--------------- core/scanner.py | 34 ++++++++++++++-- core/sources/civitai.py | 40 +++++++++++++++++-- core/workflow_analyzer.py | 15 ++++--- web/linker.js | 4 ++ 7 files changed, 177 insertions(+), 49 deletions(-) diff --git a/__init__.py b/__init__.py index 7283148..2073ef0 100644 --- a/__init__.py +++ b/__init__.py @@ -1040,6 +1040,16 @@ class ModelLinkerExtension: } category = category_map.get(raw_category.lower(), raw_category.lower()) + if not category or category == "unknown": + return web.json_response([]) + + known_categories = set(folder_paths.folder_names_and_paths.keys()) + if category not in known_categories: + self.logger.debug( + f"Model Linker: skipping subfolder lookup for unknown category '{raw_category}' -> '{category}'" + ) + return web.json_response([]) + subfolders = set() filenames = folder_paths.get_filename_list(category) or [] for rel_path in filenames: diff --git a/core/linker.py b/core/linker.py index 4abb710..3529a66 100644 --- a/core/linker.py +++ b/core/linker.py @@ -269,10 +269,21 @@ def analyze_and_find_matches( log_debug(f"Extracted {len(workflow_urls)} URLs from workflow") # Analyze workflow to find all model references - all_model_refs = analyze_workflow_models(workflow_json) - # Get available models available_models = get_model_files() + available_models_by_category = {} + for model in available_models: + model_category = model.get("category", "") + if model_category not in available_models_by_category: + available_models_by_category[model_category] = [] + available_models_by_category[model_category].append(model) + + ordered_candidates_cache: Dict[str, List[Dict[str, Any]]] = {} + + # Analyze workflow using the same already-scanned model list + all_model_refs = analyze_workflow_models( + workflow_json, available_models=available_models + ) # Identify missing models missing_models = identify_missing_models(all_model_refs, available_models) @@ -317,6 +328,19 @@ def analyze_and_find_matches( # For URNs, prefer expected_filename for matching if missing.get("is_urn") and missing.get("expected_filename"): target_for_matching = missing["expected_filename"] + elif isinstance(target_for_matching, str) and target_for_matching.startswith( + "urn:air:" + ): + # Don't fuzzy-match the full URN string against every local model. + # For worker-asset style URNs, try using the filename-like suffix after "@"; + # otherwise skip local matching until the frontend resolves the URN async. + urn_suffix = target_for_matching.split("@", 1)[1] if "@" in target_for_matching else "" + urn_suffix_ext = os.path.splitext(urn_suffix)[1].lower() + if urn_suffix and urn_suffix_ext in MODEL_EXTENSIONS: + target_for_matching = urn_suffix + else: + missing_with_matches.append({**missing, "matches": []}) + continue # Filter available models by category if known category = missing.get("category") @@ -328,10 +352,14 @@ def analyze_and_find_matches( candidates = available_models if category and category != "unknown": - candidates = [m for m in available_models if m.get("category") == category] - candidates.extend( - [m for m in available_models if m.get("category") != category] - ) + candidates = ordered_candidates_cache.get(category) + if candidates is None: + preferred = available_models_by_category.get(category, []) + others = [ + m for m in available_models if m.get("category") != category + ] + candidates = preferred + others + ordered_candidates_cache[category] = candidates # Find matches matches = find_matches( diff --git a/core/matcher.py b/core/matcher.py index 108c365..b551a60 100644 --- a/core/matcher.py +++ b/core/matcher.py @@ -6,6 +6,7 @@ Implements fuzzy string matching to find similar model names. import os import re +import heapq from typing import List, Dict, Tuple from difflib import SequenceMatcher @@ -98,7 +99,9 @@ def find_matches( 'confidence': confidence percentage (0 to 100) } """ - matches = [] + # Keep only the best N matches instead of collecting and sorting everything. + best_matches = [] + match_counter = 0 # Normalize path separators in target_model based on current OS # This ensures paths with \ vs / separators are treated as identical @@ -110,6 +113,7 @@ def find_matches( # Normalize target filename once for exact match comparisons target_norm = normalize_filename(target_filename) + target_base = os.path.splitext(target_filename)[0] for candidate in candidate_models: # Get filename from candidate (prefer 'filename' key, fallback to extracting from 'path' or 'relative_path') @@ -126,13 +130,20 @@ def find_matches( # Normalize candidate path separators based on current OS # This ensures paths with \ vs / separators are treated as identical - candidate_path_normalized = ( - os.path.normpath(candidate_path) if candidate_path else "" - ) + candidate_path_normalized = candidate.get("_match_path_norm") + if candidate_path_normalized is None: + candidate_path_normalized = ( + os.path.normpath(candidate_path) if candidate_path else "" + ) + candidate["_match_path_norm"] = candidate_path_normalized + candidate_relative_path = candidate.get("relative_path", "") - candidate_relative_path_normalized = ( - os.path.normpath(candidate_relative_path) if candidate_relative_path else "" - ) + candidate_relative_path_normalized = candidate.get("_match_relative_path_norm") + if candidate_relative_path_normalized is None: + candidate_relative_path_normalized = ( + os.path.normpath(candidate_relative_path) if candidate_relative_path else "" + ) + candidate["_match_relative_path_norm"] = candidate_relative_path_normalized # Check if normalized paths are identical (100% match) # This handles cases where paths differ only by separator (e.g., path/to/model vs path\to\model) @@ -149,14 +160,20 @@ def find_matches( if path_match: # Exact path match after normalization = 100% confidence similarity = 1.0 - matches.append( - { - "model": candidate, - "filename": candidate_filename, - "similarity": similarity, - "confidence": round(similarity * 100, 1), - } - ) + match = { + "model": candidate, + "filename": candidate_filename, + "similarity": similarity, + "confidence": round(similarity * 100, 1), + } + if max_results <= 0: + return [] + entry = (similarity, match_counter, match) + match_counter += 1 + if len(best_matches) < max_results: + heapq.heappush(best_matches, entry) + elif similarity > best_matches[0][0]: + heapq.heapreplace(best_matches, entry) continue # Calculate similarity comparing just filenames (not paths) @@ -164,7 +181,10 @@ def find_matches( # First check for exact match (after normalization) - should be 100% # Only exact matches should get 100% confidence - candidate_norm = normalize_filename(candidate_filename) + candidate_norm = candidate.get("_match_filename_norm") + if candidate_norm is None: + candidate_norm = normalize_filename(candidate_filename) + candidate["_match_filename_norm"] = candidate_norm if target_norm == candidate_norm: # Exact match after normalization = 100% confidence @@ -177,7 +197,6 @@ def find_matches( ) # Also try comparing without extensions for better matching - target_base = os.path.splitext(target_filename)[0] candidate_base = os.path.splitext(candidate_filename)[0] similarity_no_ext = calculate_similarity_with_normalization( target_base, candidate_base @@ -195,19 +214,19 @@ def find_matches( # Only include if above threshold if similarity >= threshold: - matches.append( - { - "model": candidate, - "filename": candidate_filename, - "similarity": similarity, - "confidence": round(similarity * 100, 1), # Convert to percentage - } - ) + match = { + "model": candidate, + "filename": candidate_filename, + "similarity": similarity, + "confidence": round(similarity * 100, 1), # Convert to percentage + } + if max_results <= 0: + return [] + entry = (similarity, match_counter, match) + match_counter += 1 + if len(best_matches) < max_results: + heapq.heappush(best_matches, entry) + elif similarity > best_matches[0][0]: + heapq.heapreplace(best_matches, entry) - # Sort by similarity (highest first) - matches.sort(key=lambda x: x["similarity"], reverse=True) - - # Limit to max_results - matches = matches[:max_results] - - return matches + return [match for _, _, match in sorted(best_matches, key=lambda x: x[0], reverse=True)] diff --git a/core/scanner.py b/core/scanner.py index ab52e43..5bdfcf9 100644 --- a/core/scanner.py +++ b/core/scanner.py @@ -5,7 +5,8 @@ Scans configured model directories and finds available model files. """ import os -from typing import List, Dict, Tuple +import time +from typing import List, Dict, Tuple, Optional from .log_system.log_funcs import ( log_debug, @@ -37,6 +38,10 @@ MODEL_EXTENSIONS = { ".gguf", } +_MODEL_FILES_CACHE: Optional[List[Dict[str, str]]] = None +_MODEL_FILES_CACHE_AT: float = 0.0 +_MODEL_FILES_CACHE_TTL_SECONDS = 2.0 + def get_model_directories() -> Dict[str, Tuple[List[str], set]]: """ @@ -199,13 +204,36 @@ def scan_all_directories() -> List[Dict[str, str]]: return all_models -def get_model_files() -> List[Dict[str, str]]: +def invalidate_model_files_cache() -> None: + """Clear the in-memory model file cache.""" + global _MODEL_FILES_CACHE, _MODEL_FILES_CACHE_AT + _MODEL_FILES_CACHE = None + _MODEL_FILES_CACHE_AT = 0.0 + + +def get_model_files(force_rescan: bool = False) -> List[Dict[str, str]]: """ Get list of all available model files with metadata. This is the main entry point for getting model files. + Args: + force_rescan: If True, bypass the short-lived cache and rescan directories + Returns: List of model dictionaries (same format as scan_directory) """ - return scan_all_directories() + global _MODEL_FILES_CACHE, _MODEL_FILES_CACHE_AT + + now = time.monotonic() + if ( + not force_rescan + and _MODEL_FILES_CACHE is not None + and (now - _MODEL_FILES_CACHE_AT) < _MODEL_FILES_CACHE_TTL_SECONDS + ): + return _MODEL_FILES_CACHE + + models = scan_all_directories() + _MODEL_FILES_CACHE = models + _MODEL_FILES_CACHE_AT = now + return models diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 9679dd7..1c6c798 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -28,6 +28,41 @@ _urn_cache: Dict[tuple[int, int], Dict[str, Any]] = {} _hash_cache: Dict[str, Dict[str, Any]] = {} +def _extract_civitai_image_id(image_url: str) -> Optional[str]: + """ + Extract a CivitAI image ID from an image CDN URL. + Example: + https://image.civitai.com/.../width=1800/1917130.jpeg -> 1917130 + """ + if not image_url: + return None + + match = re.search(r"/(\d+)(?:\.[A-Za-z0-9]+)?(?:[?#].*)?$", image_url) + if match: + return match.group(1) + + return None + + +def _build_civitai_image_url(img: Dict[str, Any]) -> str: + """ + Build a stable CivitAI image page URL from available image metadata. + """ + civitai_url = img.get("civitaiUrl") + if civitai_url: + return civitai_url + + image_id = img.get("id") + if image_id is not None: + return f"https://civitai.com/images/{image_id}" + + extracted_id = _extract_civitai_image_id(img.get("url", "")) + if extracted_id: + return f"https://civitai.com/images/{extracted_id}" + + return "" + + def parse_civitai_url(url: str) -> Optional[Dict[str, Any]]: """ Parse a CivitAI URL to extract model/version info. @@ -650,7 +685,7 @@ def _extract_model_images(version_info: Dict[str, Any]) -> List[Dict[str, Any]]: img_info = { "url": img_url, - "civitaiUrl": img.get("civitaiUrl", ""), + "civitaiUrl": _build_civitai_image_url(img), "seed": img.get("seed") or meta.get("seed"), "steps": img.get("steps") or meta.get("steps"), "cfg": img.get("cfg") or meta.get("cfg"), @@ -743,8 +778,7 @@ def _metadata_to_model_info(metadata: Dict[str, Any]) -> Dict[str, Any]: images.append( { "url": img.get("url", ""), - "civitaiUrl": img.get("civitaiUrl", "") - or f"https://civitai.com/images/{img.get('id')}", + "civitaiUrl": _build_civitai_image_url(img), "seed": img_meta.get("seed"), "steps": img_meta.get("steps"), "cfg": img_meta.get("cfg"), diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py index 295d344..d632b76 100644 --- a/core/workflow_analyzer.py +++ b/core/workflow_analyzer.py @@ -166,7 +166,9 @@ def try_resolve_model_path( return None -def get_node_model_info(node: Dict[str, Any]) -> List[Dict[str, Any]]: +def get_node_model_info( + node: Dict[str, Any], available_models: Optional[List[Dict[str, Any]]] = None +) -> List[Dict[str, Any]]: """ Extract model references from a single node. @@ -225,7 +227,7 @@ def get_node_model_info(node: Dict[str, Any]) -> List[Dict[str, Any]]: # Get all lora files using scanner for recursive search from .scanner import get_model_files - all_loras = get_model_files() + all_loras = available_models if available_models is not None else get_model_files() lora_files = [m for m in all_loras if m.get("category") == "loras"] # Build a lookup by filename (without extension) @@ -407,7 +409,10 @@ def get_node_model_info(node: Dict[str, Any]) -> List[Dict[str, Any]]: return model_refs -def analyze_workflow_models(workflow_json: Dict[str, Any]) -> List[Dict[str, Any]]: +def analyze_workflow_models( + workflow_json: Dict[str, Any], + available_models: Optional[List[Dict[str, Any]]] = None, +) -> List[Dict[str, Any]]: """ Extract all model references from a workflow, including nested subgraphs. @@ -429,7 +434,7 @@ def analyze_workflow_models(workflow_json: Dict[str, Any]) -> List[Dict[str, Any nodes = workflow_json.get("nodes", []) for node in nodes: try: - model_refs = get_node_model_info(node) + model_refs = get_node_model_info(node, available_models=available_models) node_type = node.get("type", "") # Check if node type is a subgraph UUID @@ -467,7 +472,7 @@ def analyze_workflow_models(workflow_json: Dict[str, Any]) -> List[Dict[str, Any for node in subgraph_nodes: try: - model_refs = get_node_model_info(node) + model_refs = get_node_model_info(node, available_models=available_models) # Mark as belonging to this subgraph definition for ref in model_refs: ref["subgraph_id"] = subgraph_id diff --git a/web/linker.js b/web/linker.js index b9c9e6a..2c630bc 100644 --- a/web/linker.js +++ b/web/linker.js @@ -2504,6 +2504,10 @@ class LinkerManagerDialog extends ComfyDialog { async ensureDownloadSubfoldersLoaded(category = '') { const key = (category || '').trim().toLowerCase(); if (!key) return []; + if (key === 'unknown') { + this.downloadSubfolders.set(key, []); + return []; + } if (this.downloadSubfolders.has(key)) { return this.downloadSubfolders.get(key) || []; }