""" Core Resolver Module Integrates all components to provide high-level API for model linking. """ import os import re import json from typing import Dict, Any, List, Optional, Tuple, Callable from urllib.parse import unquote from .log_system.log_funcs import ( log_debug, log_info, log_warn, log_error, log_exception, ) from .scanner import get_model_files from .workflow_analyzer import analyze_workflow_models, identify_missing_models from .matcher import find_matches from .workflow_updater import update_workflow_nodes from .sources.civitai import resolve_urn # Regex patterns for URL extraction (matches HuggingFace and CivitAI URLs) URL_PATTERN = re.compile(r'(https?://(?:huggingface\.co|civitai\.com)[^\s"\'<>\)\\]+)') # Model file extensions to look for MODEL_EXTENSIONS = (".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".onnx", ".gguf") def strip_known_model_extension(filename: str) -> str: """Strip only known model extensions, preserving names like v4.0.""" if not isinstance(filename, str): return "" lowered = filename.lower() for ext in MODEL_EXTENSIONS: if lowered.endswith(ext): return filename[: -len(ext)] return filename def get_workflow_url_info_for_filename( workflow_urls: Dict[str, Dict[str, Any]], filename: str ) -> Optional[Dict[str, Any]]: if filename in workflow_urls: return workflow_urls[filename] filename_stem = strip_known_model_extension(os.path.basename(filename)).lower() if not filename_stem: return None for workflow_filename, url_info in workflow_urls.items(): workflow_stem = strip_known_model_extension( os.path.basename(workflow_filename) ).lower() if workflow_stem == filename_stem: return url_info return None def normalize_workflow_download_url(url: str) -> str: """Convert workflow file-page URLs into direct download URLs when possible.""" if not isinstance(url, str) or not url: return url # HuggingFace /blob/ pages return HTML. /resolve/ returns the actual file. return re.sub( r"(https?://huggingface\.co/[^/]+/[^/]+)/blob/([^/]+)/(.+)", r"\1/resolve/\2/\3", url, ) def workflow_url_points_to_file(url: str, filename: str) -> bool: """Return true when a URL appears to reference the specific model file.""" if not url or not filename: return False try: decoded_url = unquote(url) except Exception: decoded_url = url return filename in decoded_url or unquote(filename) in decoded_url def search_local_matches( target_for_matching: str, category: Optional[str] = None, similarity_threshold: float = 0.0, max_matches_per_model: int = 10, force_rescan: bool = False, ) -> List[Dict[str, Any]]: """ Search local model files using the same matcher as workflow analysis. Args: target_for_matching: Filename/path to match against local files category: Optional category hint to prioritize/filter candidates similarity_threshold: Minimum similarity score (0.0 to 1.0) max_matches_per_model: Maximum number of matches to return Returns: Deduplicated list of local matches sorted by similarity """ available_models = get_model_files(force_rescan=force_rescan) 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] ) matches = find_matches( target_for_matching, candidates, threshold=similarity_threshold, max_results=max_matches_per_model, ) seen_absolute_paths = {} deduplicated_matches = [] for match in matches: model_dict = match["model"] absolute_path = model_dict.get("path", "") if absolute_path: absolute_path = os.path.normpath(absolute_path) if absolute_path not in seen_absolute_paths: seen_absolute_paths[absolute_path] = match deduplicated_matches.append(match) else: existing_match = seen_absolute_paths[absolute_path] if match["confidence"] > existing_match["confidence"]: idx = deduplicated_matches.index(existing_match) deduplicated_matches[idx] = match seen_absolute_paths[absolute_path] = match return deduplicated_matches def extract_workflow_urls(workflow_json: Dict[str, Any]) -> Dict[str, Dict[str, Any]]: """ Extract model URLs from workflow JSON. Sources: 1. node.properties.models array - contains {name, url, directory} 2. Regex extraction from workflow JSON string - finds HuggingFace/CivitAI URLs Args: workflow_json: Complete workflow JSON dictionary Returns: Dict mapping model filename -> {url, directory, source} """ url_map = {} # Convert to string for regex search workflow_str = json.dumps(workflow_json) # Collect all nodes including from subgraphs all_nodes = list(workflow_json.get("nodes", [])) definitions = workflow_json.get("definitions", {}) subgraphs = definitions.get("subgraphs", []) for subgraph in subgraphs: subgraph_nodes = subgraph.get("nodes", []) all_nodes.extend(subgraph_nodes) # 1. Extract from node.properties.models (authoritative source) for node in all_nodes: node_type = node.get("type", "") properties = node.get("properties", {}) models_list = properties.get("models", []) for model_info in models_list: if isinstance(model_info, dict): name = model_info.get("name", "") url = model_info.get("url", "") directory = model_info.get("directory", "") if name and name not in url_map: url_map[name] = { "url": normalize_workflow_download_url(url), "model_url": url, "directory": directory, "node_type": node_type, "source": "node_properties", } # 2. Extract URLs via regex from workflow JSON urls_found = URL_PATTERN.findall(workflow_str) # Clean URLs (remove trailing characters that may have been captured) cleaned_urls = [] for url in urls_found: url = url.split(")")[0].replace("\\n", "").replace("\n", "").strip() if url: cleaned_urls.append(url) # 3. Extract model filenames via regex model_pattern = re.compile( r"([\w\-\.%]+\.(?:safetensors|ckpt|pt|pth|bin|onnx|gguf))", re.IGNORECASE ) model_files_raw = model_pattern.findall(workflow_str) # Clean and decode filenames model_files = set() model_name_map = {} # decoded -> original for model in model_files_raw: cleaned = model.strip() if cleaned and cleaned[0].isalnum(): try: decoded = unquote(cleaned) except Exception: decoded = cleaned model_files.add(decoded) model_name_map[decoded] = cleaned # 4. Match URLs to model filenames for model in model_files: # Keep authoritative node.properties URLs only when they point to the file. # Some workflows store a repo/model-page URL there and a concrete file URL # elsewhere in the JSON; in that case the file URL is the usable download. if ( model in url_map and url_map[model].get("url") and workflow_url_points_to_file(url_map[model].get("url"), model) ): continue original_name = model_name_map.get(model, model) for url in cleaned_urls: # Check decoded name in URL if model in url: if ( model not in url_map or not url_map[model].get("url") or not workflow_url_points_to_file(url_map[model].get("url"), model) ): url_map[model] = { "url": normalize_workflow_download_url(url), "model_url": url, "directory": url_map.get(model, {}).get("directory", ""), "source": "regex", } else: url_map[model]["url"] = normalize_workflow_download_url(url) url_map[model]["model_url"] = url url_map[model]["source"] = "regex" break # Check original (possibly URL-encoded) name in URL if original_name in url: if ( model not in url_map or not url_map[model].get("url") or not workflow_url_points_to_file( url_map[model].get("url"), original_name ) ): url_map[model] = { "url": normalize_workflow_download_url(url), "model_url": url, "directory": url_map.get(model, {}).get("directory", ""), "source": "regex", } else: url_map[model]["url"] = normalize_workflow_download_url(url) url_map[model]["model_url"] = url url_map[model]["source"] = "regex" break # Check without extension model_base = os.path.splitext(model)[0] if model_base in url or unquote(model_base) in url: if ( model not in url_map or not url_map[model].get("url") or not workflow_url_points_to_file(url_map[model].get("url"), model) ): url_map[model] = { "url": normalize_workflow_download_url(url), "model_url": url, "directory": url_map.get(model, {}).get("directory", ""), "source": "regex", } else: url_map[model]["url"] = normalize_workflow_download_url(url) url_map[model]["model_url"] = url url_map[model]["source"] = "regex" break return url_map def parse_huggingface_url(url: str) -> Tuple[Optional[str], Optional[str]]: """ Extract HuggingFace repo and path from URL. Args: url: HuggingFace URL Returns: Tuple of (repo_id, file_path) or (None, None) if not valid """ if not url or "huggingface.co" not in url: return None, None # Pattern: https://huggingface.co/user/repo/resolve/main/path/to/file.safetensors match = re.match( r"https?://huggingface\.co/([^/]+/[^/]+)/(?:resolve|blob)/[^/]+/(.+)", url ) if match: return match.group(1), match.group(2) return None, None def analyze_and_find_matches( workflow_json: Dict[str, Any], similarity_threshold: float = 0.0, max_matches_per_model: int = 10, progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None, force_rescan: bool = False, ) -> Dict[str, Any]: """ Main entry point: analyze workflow and find matches for missing models. Args: workflow_json: Complete workflow JSON dictionary similarity_threshold: Minimum similarity score (0.0 to 1.0) for matches max_matches_per_model: Maximum number of matches to return per missing model force_rescan: If True, bypass the short-lived local model scan cache Returns: Dictionary with analysis results: { 'missing_models': [ { 'node_id': node ID, 'node_type': node type, 'widget_index': widget index, 'original_path': original path from workflow, 'category': model category, 'workflow_url': URL from workflow if found, 'workflow_directory': directory from workflow if found, 'matches': [ { 'model': model dict from scanner, 'filename': model filename, 'similarity': similarity score (0.0-1.0), 'confidence': confidence percentage (0-100) }, ... ] }, ... ], 'total_missing': count of missing models, 'total_models_analyzed': count of all models in workflow } """ if progress_callback: progress_callback( { "stage": "extracting", "message": "Extracting workflow model references...", "current": 0, "total": 0, } ) # Extract URLs from workflow (node.properties.models + regex) workflow_urls = extract_workflow_urls(workflow_json) log_debug(f"Extracted {len(workflow_urls)} URLs from workflow") if progress_callback: progress_callback( { "stage": "scanning", "message": "Scanning local model index...", "current": 0, "total": 0, } ) # Analyze workflow to find all model references # Get available models available_models = get_model_files(force_rescan=force_rescan) 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]]] = {} if progress_callback: progress_callback( { "stage": "analyzing", "message": "Analyzing workflow nodes...", "current": 0, "total": 0, } ) # Analyze workflow using the same already-scanned model list all_model_refs = analyze_workflow_models( workflow_json, available_models=available_models ) if progress_callback: progress_callback( { "stage": "identifying", "message": "Identifying missing models...", "current": 0, "total": len(all_model_refs), } ) # Identify missing models missing_models = identify_missing_models(all_model_refs, available_models) # Enrich missing models with workflow URLs for missing in missing_models: original_path = missing.get("original_path", "") filename = os.path.basename(original_path) url_info = get_workflow_url_info_for_filename(workflow_urls, filename) if url_info: missing["workflow_url"] = url_info.get("url", "") missing["workflow_model_url"] = url_info.get("model_url", "") missing["workflow_directory"] = url_info.get("directory", "") missing["url_source"] = url_info.get("source", "") # Handle URNs: mark for async resolution by frontend # No sync CivitAI calls here - frontend will fetch asynchronously for missing in missing_models: if missing.get("is_urn"): missing["needs_urn_resolve"] = True urn = missing.get("urn") if urn: missing["urn_model_id"] = urn.get("model_id") missing["urn_version_id"] = urn.get("version_id") missing["urn_type"] = urn.get("type", "") if progress_callback: progress_callback( { "stage": "matching", "message": "Matching missing models...", "current": 0, "total": len(missing_models), } ) # Find matches for each missing model missing_with_matches = [] total_missing = len(missing_models) for index, missing in enumerate(missing_models, start=1): if progress_callback: progress_callback( { "stage": "matching", "message": f"Analyzing model {index} of {total_missing}", "current": index, "total": total_missing, "model_name": missing.get("name") or missing.get("original_path", ""), } ) # Skip LoraManager lorAs that already exist locally (exists=True means no linking needed) is_lora_v2 = missing.get("is_lora_v2") exists = missing.get("exists") name = missing.get("name") or missing.get("original_path", "") log_debug(f"Checking {name}: is_lora_v2={is_lora_v2}, exists={exists}") if is_lora_v2 and exists: log_info(f"Skipping LoraManager lora {name} - already exists locally") continue target_for_matching = missing.get("original_path", "") # 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") if not category or category == "unknown": from .workflow_analyzer import NODE_TYPE_TO_CATEGORY_HINTS node_type = missing.get("node_type", "") category = NODE_TYPE_TO_CATEGORY_HINTS.get(node_type, "unknown") candidates = available_models if category and category != "unknown": 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( target_for_matching, candidates, threshold=similarity_threshold, max_results=max_matches_per_model, ) # Deduplicate matches by absolute path seen_absolute_paths = {} deduplicated_matches = [] for match in matches: model_dict = match["model"] absolute_path = model_dict.get("path", "") if absolute_path: absolute_path = os.path.normpath(absolute_path) if absolute_path not in seen_absolute_paths: seen_absolute_paths[absolute_path] = match deduplicated_matches.append(match) else: existing_match = seen_absolute_paths[absolute_path] if match["confidence"] > existing_match["confidence"]: idx = deduplicated_matches.index(existing_match) deduplicated_matches[idx] = match seen_absolute_paths[absolute_path] = match missing_with_matches.append({**missing, "matches": deduplicated_matches}) result = { "missing_models": missing_with_matches, "total_missing": len(missing_with_matches), "total_models_analyzed": len(all_model_refs), } if progress_callback: progress_callback( { "stage": "completed", "message": "Analysis complete", "current": total_missing, "total": total_missing, } ) return result def apply_resolution( workflow_json: Dict[str, Any], resolutions: List[Dict[str, Any]] ) -> Dict[str, Any]: """ Apply model resolutions to workflow. Args: workflow_json: Workflow JSON dictionary (will be modified) resolutions: List of resolution dictionaries: { 'node_id': node ID, 'widget_index': widget index, 'resolved_path': absolute path to resolved model, 'category': model category (optional), 'resolved_model': model dict from scanner (optional), 'nested_key': nested key for dict-type widgets (optional) } Returns: Updated workflow JSON dictionary """ # Prepare mappings for workflow_updater mappings = [] for resolution in resolutions: mapping = { "node_id": resolution.get("node_id"), "widget_index": resolution.get("widget_index"), "resolved_path": resolution.get("resolved_path"), "category": resolution.get("category"), "resolved_model": resolution.get("resolved_model"), "subgraph_id": resolution.get( "subgraph_id" ), # Include subgraph_id for subgraph nodes "is_top_level": resolution.get( "is_top_level" ), # True for top-level nodes, False for nodes in subgraph definitions "is_lora_v2": resolution.get("is_lora_v2"), # Flag for LoraManager nodes "original_lora_name": resolution.get( "original_lora_name" ), # Original lora name for LoraManager replacement "nested_key": resolution.get( "nested_key" ), # For dict-type widgets (e.g. Power Lora Loader) } # If resolved_model provided, extract path if needed if "resolved_model" in resolution and resolution["resolved_model"]: resolved_model = resolution["resolved_model"] if "path" in resolved_model and not mapping.get("resolved_path"): mapping["resolved_path"] = resolved_model["path"] if "base_directory" in resolved_model: mapping["base_directory"] = resolved_model["base_directory"] mappings.append(mapping) # Update workflow updated_workflow = update_workflow_nodes(workflow_json, mappings) return updated_workflow def get_resolution_summary(workflow_json: Dict[str, Any]) -> Dict[str, Any]: """ Get summary of missing models and matches without applying resolutions. This is a convenience method that calls analyze_and_find_matches with defaults. Args: workflow_json: Complete workflow JSON dictionary Returns: Same format as analyze_and_find_matches """ return analyze_and_find_matches(workflow_json)