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.
2026-04-24 23:32:00 +02:00
parent 2b177817f4
commit e096f2c786
7 changed files with 177 additions and 49 deletions
@@ -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:
@@ -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(
@@ -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)]
@@ -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
@@ -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"),
@@ -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
@@ -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) || [];
        }