From e3fdf9f3345ca6f71d005966c5c2033d8de7cc6e Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Sun, 2 Aug 2026 19:28:33 +0200 Subject: [PATCH] refactor: extract workflow references --- core/resolver.py | 9 +- core/workflow/references.py | 649 ++++++++++++++++++++++++++++++ core/workflow_analyzer.py | 688 +------------------------------- tests/test_workflow_analyzer.py | 22 +- 4 files changed, 670 insertions(+), 698 deletions(-) create mode 100644 core/workflow/references.py diff --git a/core/resolver.py b/core/resolver.py index a7f76df..e3a1320 100644 --- a/core/resolver.py +++ b/core/resolver.py @@ -31,13 +31,14 @@ from .type_utils import ( prepare_remote_size_probe_url, unique_ordered_strings, ) -from .workflow_analyzer import ( - NESTED_MODEL_KEYS, - NODE_TYPE_TO_CATEGORY_HINTS, +from .workflow.references import ( get_model_widget_category_hint, + should_scan_as_model_reference, +) +from .workflow.widgets import NESTED_MODEL_KEYS, NODE_TYPE_TO_CATEGORY_HINTS +from .workflow_analyzer import ( get_workflow_model_inventory, identify_missing_models, - should_scan_as_model_reference, ) from .workflow_updater import update_workflow_nodes diff --git a/core/workflow/references.py b/core/workflow/references.py new file mode 100644 index 0000000..cdfb59f --- /dev/null +++ b/core/workflow/references.py @@ -0,0 +1,649 @@ +"""Model reference matching, category inference, and path resolution.""" + +import os +import re +from typing import Any, Dict, List, Optional + +from ..custom_nodes import ( + analyze_custom_node_references, + get_custom_node_model_adapter, +) +from ..log_system import create_module_logger +from ..type_utils import MODEL_EXTENSIONS, URN_TYPE_MAP, normalize_download_category +from . import dynamic_widgets +from .widgets import ( + NESTED_MODEL_KEYS, + NODE_TYPE_TO_CATEGORY_HINTS, + _ordered_unique_categories, + get_node_model_widget_category_hint, + get_node_output_category_hint, + get_widget_category_hint, + get_widget_name_hint, + is_workflow_model_widget_candidate, +) + +log = create_module_logger(__name__) + +try: + import folder_paths +except ImportError: + folder_paths = None + + +URN_REGEX = re.compile(r"^urn:air:([^:]+):([^:]+):([^:]+):(\d+)@(\d+)$") + + +def is_static_or_hybrid_widget_choice(value: Any, choice_info: Dict[str, Any]) -> bool: + source = str(choice_info.get("source") or "").strip().lower() + if source not in {"static", "hybrid"}: + return False + + target = dynamic_widgets._normalize_choice_for_match(value) + if not target: + return False + + choices = choice_info.get("choices") or [] + return any( + dynamic_widgets._normalize_choice_for_match(choice) == target + for choice in choices + ) + + +def static_or_hybrid_choice_looks_like_model( + value: Any, choice_info: Dict[str, Any] +) -> bool: + if is_model_filename(value): + return True + + if not isinstance(value, str): + return False + + value_text = value.strip() + if not value_text: + return False + + if "/" in value_text or "\\" in value_text: + return True + + target = dynamic_widgets._normalize_choice_for_match(value_text) + choices = choice_info.get("choices") or [] + return any( + dynamic_widgets._normalize_choice_for_match(choice) == target + and is_model_filename(choice) + for choice in choices + ) + + +MODEL_WIDGET_PLACEHOLDERS = { + "", + ".none", + ".use_ckpt_clip", + ".use_ckpt_vae", + "none", + "[none]", + "null", + "undefined", + "default", + "auto", + "baked vae", + "included", + "(use same)", + "select the lora to add to the text", + "esam", + "pixel_space", + "taesd", + "taesdxl", + "taesd3", + "taef1", + "taef2", +} + + +def get_model_widget_category_hint( + node: Dict[str, Any], widget_index: int +) -> Optional[str]: + category_hints = get_model_widget_category_hints(node, widget_index) + return category_hints[0] if category_hints else None + + +def get_model_widget_category_hints( + node: Dict[str, Any], widget_index: int +) -> List[str]: + node_type = node.get("type", "") + indexed_category_hint = get_node_model_widget_category_hint(node_type, widget_index) + widget_category_hint = get_widget_category_hint(node, widget_index) + dynamic_category_hints = dynamic_widgets.get_dynamic_widget_category_hints( + node, widget_index + ) + + output_category_hint = get_node_output_category_hint(node) + output_widget_category_hint = ( + output_category_hint + if ( + not dynamic_category_hints + and (indexed_category_hint or widget_category_hint) + ) + else None + ) + + hints_list = [] + for hint in dynamic_category_hints: + if hint: + hints_list.append(normalize_download_category(hint)) + + if dynamic_category_hints: + return _ordered_unique_categories(hints_list) + + if indexed_category_hint: + hints_list.append(indexed_category_hint) + + if ( + widget_category_hint == "diffusion_models" + and output_widget_category_hint + and output_widget_category_hint != "diffusion_models" + ): + hints_list.append(output_widget_category_hint) + hints_list.append(widget_category_hint) + else: + if widget_category_hint: + hints_list.append(widget_category_hint) + if output_widget_category_hint: + hints_list.append(output_widget_category_hint) + + return _ordered_unique_categories(hints_list) + + +def get_effective_model_category_hint( + node: Dict[str, Any], widget_index: int +) -> Optional[str]: + return get_model_widget_category_hint( + node, widget_index + ) or NODE_TYPE_TO_CATEGORY_HINTS.get(node.get("type", "")) + + +def is_placeholder_model_value(value: Any) -> bool: + if not isinstance(value, str): + return False + + return value.strip().lower() in MODEL_WIDGET_PLACEHOLDERS + + +def should_scan_as_model_reference(value: Any, declared_model_widget: bool) -> bool: + """Detect references only when the value or its widget provides model evidence.""" + if isinstance(value, str) and URN_REGEX.match(value.strip()): + return True + + if not declared_model_widget or not isinstance(value, str): + return False + + if is_model_filename(value): + return True + + return bool(value.strip()) and not is_placeholder_model_value(value) + + +def is_model_filename(value: Any) -> bool: + """Return whether a value looks like a model filename or URN.""" + if not isinstance(value, str): + return False + + _, extension = os.path.splitext(value.lower()) + if extension in MODEL_EXTENSIONS: + return True + + return bool(URN_REGEX.match(value.strip())) + + +def _normalize_model_path_for_lookup(value: str) -> str: + """Normalize separators for comparison while preserving letter case.""" + if not isinstance(value, str): + return "" + + normalized = os.path.normpath(value.strip()) + if normalized == ".": + return "" + + return normalized.replace("\\", "/").strip("/") + + +def _resolve_from_available_models( + filename: str, + categories: Optional[List[str]], + available_models: Optional[List[Dict[str, Any]]], +) -> Optional[tuple[str, str]]: + """Resolve using scanner data with case-sensitive relative path matching.""" + if not available_models: + return None + + requested_key = _normalize_model_path_for_lookup(filename) + if not requested_key: + return None + + requested_is_absolute = os.path.isabs(filename) + if requested_is_absolute: + requested_key = _normalize_model_path_for_lookup(os.path.abspath(filename)) + + if categories is None: + category_order = [] + seen_categories = set() + for model in available_models: + category = model.get("category") + if category and category not in seen_categories: + seen_categories.add(category) + category_order.append(category) + else: + category_order = categories + + if not category_order: + return None + + for category in category_order: + normalized_category = normalize_download_category(category) + for model in available_models: + model_category = model.get("category") + if ( + model_category != category + and normalize_download_category(model_category) != normalized_category + ): + continue + + if requested_is_absolute: + model_path = model.get("path") or "" + else: + model_path = model.get("relative_path") or model.get("filename") or "" + + if _normalize_model_path_for_lookup(model_path) != requested_key: + continue + + full_path = model.get("path") + if full_path and os.path.exists(full_path): + return (category, full_path) + + return None + + +def _category_has_exact_filename(category: str, filename: str) -> Optional[bool]: + """Check ComfyUI's filename list for an exact-case relative path match.""" + if os.path.isabs(filename): + return None + + requested_key = _normalize_model_path_for_lookup(filename) + if not requested_key: + return False + + try: + available_filenames = folder_paths.get_filename_list(category) or [] + except Exception: + return None + + return any( + _normalize_model_path_for_lookup(available_filename) == requested_key + for available_filename in available_filenames + if isinstance(available_filename, str) + ) + + +def try_resolve_model_path( + value: str, + categories: List[str] = None, + available_models: Optional[List[Dict[str, Any]]] = None, +) -> Optional[tuple[str, str]]: + """Resolve a workflow model value using scanner data or ComfyUI paths.""" + if not isinstance(value, str) or not value.strip(): + return None + + filename = value.strip() + + if categories is not None: + skip_categories = {"custom_nodes", "configs"} + categories = [category for category in categories if category not in skip_categories] + + resolved = _resolve_from_available_models(filename, categories, available_models) + if resolved: + return resolved + + if available_models is not None: + return None + + global folder_paths + if folder_paths is None: + try: + import folder_paths as folder_paths_module + + folder_paths = folder_paths_module + except ImportError: + log.error("Model Resolver: folder_paths not available") + return None + + if categories is None: + categories = list(folder_paths.folder_names_and_paths.keys()) + + skip_categories = {"custom_nodes", "configs"} + categories = [category for category in categories if category not in skip_categories] + + for category in categories: + try: + exact_filename = _category_has_exact_filename(category, filename) + if exact_filename is False: + continue + + full_path = folder_paths.get_full_path(category, filename) + if full_path and os.path.exists(full_path): + return (category, full_path) + except Exception: + continue + + return None + +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. + + This scans widget values backed by model metadata or a model-like widget name. + Explicit model URNs are also accepted without widget metadata. + + Args: + node: Node dictionary from workflow JSON + + Returns: + List of model reference dictionaries: + { + 'node_id': node id, + 'node_type': node type, + 'widget_index': index in widgets_values, + 'original_path': original path from workflow, + 'category': model category (if found), + 'exists': True if model exists, + 'connected': True if node has any connected inputs/outputs + } + """ + model_refs = [] + node_id = node.get("id") + node_type = node.get("type", "") + node_title = str(node.get("title", "") or "").strip() + widgets_values = node.get("widgets_values", []) + + # Check if node is connected (has any inputs or outputs with links) + inputs = node.get("inputs", []) + outputs = node.get("outputs", []) + is_connected = any(inp.get("link") is not None for inp in inputs) or any( + out.get("links") and len(out.get("links", [])) > 0 for out in outputs + ) + + # Check if node is in bypass mode (mode 4) + node_mode = node.get("mode", 0) + is_bypassed = node_mode == 4 + + # Node is active if connected AND not bypassed + is_active = is_connected and not is_bypassed + + if not widgets_values: + return model_refs + + custom_model_refs = analyze_custom_node_references( + node, + available_models, + is_active=is_active, + get_widget_name_hint=get_widget_name_hint, + ) + if custom_model_refs is not None: + return custom_model_refs + + # For each widget value, check if it looks like a model file or URN + for idx, value in enumerate(widgets_values): + if dynamic_widgets.is_dynamic_non_model_widget(node, idx): + continue + widget_name = ( + dynamic_widgets.get_dynamic_serialized_widget_name(node, idx) + or get_widget_name_hint(node, idx) + ) + dynamic_category_hints = dynamic_widgets.get_dynamic_widget_category_hints( + node, idx + ) + if dynamic_category_hints and all( + normalize_download_category(category) + in dynamic_widgets.NON_MODEL_REFERENCE_CATEGORIES + for category in dynamic_category_hints + ): + continue + model_widget_folder_key_hints = dynamic_category_hints + model_widget_choice_info = dynamic_widgets.get_dynamic_widget_choice_info( + node, idx + ) + model_widget_category_hints = get_model_widget_category_hints(node, idx) + model_widget_category_hint = ( + model_widget_category_hints[0] if model_widget_category_hints else None + ) + input_choice_source = str( + model_widget_choice_info.get("source") or "unknown" + ).strip().lower() + static_input_choice_matches_value = is_static_or_hybrid_widget_choice( + value, + model_widget_choice_info, + ) + static_input_choice_looks_like_model = static_or_hybrid_choice_looks_like_model( + value, + model_widget_choice_info, + ) + if ( + input_choice_source in {"static", "hybrid"} + and static_input_choice_matches_value + and not static_input_choice_looks_like_model + ): + continue + output_category_hint = get_node_output_category_hint(node) + workflow_schema_model_candidate = bool( + input_choice_source == "unknown" + and output_category_hint + and is_workflow_model_widget_candidate(node, idx) + and is_model_filename(value) + ) + if workflow_schema_model_candidate: + input_choice_source = "workflow_schema" + input_choice_matches_value = bool( + static_input_choice_matches_value or workflow_schema_model_candidate + ) + input_choice_looks_like_model = bool( + static_input_choice_looks_like_model or workflow_schema_model_candidate + ) + schema_output_category_hint = ( + output_category_hint + if input_choice_matches_value + and input_choice_source in {"static", "hybrid", "workflow_schema"} + and input_choice_looks_like_model + else None + ) + effective_category_hint = ( + model_widget_category_hint + or NODE_TYPE_TO_CATEGORY_HINTS.get(node_type) + or schema_output_category_hint + ) + categories_to_try_for_widget = ( + model_widget_category_hints + if model_widget_category_hints + else ([effective_category_hint] if effective_category_hint else None) + ) + # A generic widget name such as ``model_name`` is useful for choosing a + # category, but it is not proof that an arbitrary STRING value is a local + # model reference. Some metadata/search nodes use that name for human- + # readable titles. Only folder-backed schemas, known widget indexes, or a + # model-typed combo/output pairing are strong enough to admit an + # extensionless value. + category_backed_model_widget = bool( + model_widget_folder_key_hints + or get_node_model_widget_category_hint(node_type, idx) + or schema_output_category_hint + ) + named_model_file_widget = bool( + is_workflow_model_widget_candidate(node, idx) + and is_model_filename(value) + ) + + if not should_scan_as_model_reference( + value, + declared_model_widget=( + category_backed_model_widget or named_model_file_widget + ), + ): + # Check for dict-type widget values containing model references. + if isinstance(value, dict): + for nested_key, nested_category_hint in NESTED_MODEL_KEYS.items(): + nested_value = value.get(nested_key) + if ( + not nested_value + or not isinstance(nested_value, str) + or not should_scan_as_model_reference( + nested_value, declared_model_widget=True + ) + ): + continue + + value_str = nested_value.strip() + nested_categories = ( + [nested_category_hint] if nested_category_hint else None + ) + + resolved = try_resolve_model_path( + value_str, + nested_categories, + available_models=available_models, + ) + if resolved: + category, full_path = resolved + exists = os.path.exists(full_path) + else: + category = nested_category_hint or "unknown" + full_path = None + exists = False + + ref = { + "node_id": node_id, + "node_type": node_type, + "widget_index": idx, + "widget_name": widget_name, + "original_path": value_str, + "node_title": node_title, + "category": category, + "category_hints": nested_categories or ([category] if category else []), + "folder_key_hints": nested_categories or ([category] if category else []), + "full_path": full_path, + "exists": exists, + "is_urn": False, + "connected": is_active, + "nested_key": nested_key, # Track nested key for updates + } + custom_adapter = get_custom_node_model_adapter(node) + if custom_adapter: + ref["custom_node_adapter"] = custom_adapter.adapter_id + + if nested_key == "lora": + strength = value.get("strength") + if strength is not None: + try: + ref["strength"] = float(strength) + except (TypeError, ValueError): + pass + + if isinstance(value.get("on"), bool): + ref["active"] = value.get("on") + + model_refs.append(ref) + continue + + value_str = str(value).strip() + + # Check if URN + urn_match = URN_REGEX.match(value_str) + if urn_match: + base, typ, provider, model_id, version_id = urn_match.groups() + category = ( + effective_category_hint + or URN_TYPE_MAP.get(typ.lower(), "unknown") + ) + urn_category_hints = ( + categories_to_try_for_widget + if categories_to_try_for_widget + else ([category] if category else []) + ) + urn_folder_key_hints = ( + model_widget_folder_key_hints + if model_widget_folder_key_hints + else urn_category_hints + ) + + model_refs.append( + { + "node_id": node_id, + "node_type": node_type, + "widget_index": idx, + "widget_name": widget_name, + "original_path": value_str, + "urn": { + "full": value_str, + "base": base, + "type": typ, + "provider": provider, + "model_id": int(model_id), + "version_id": int(version_id), + }, + "node_title": node_title, + "category": category, + "category_hints": urn_category_hints, + "folder_key_hints": urn_folder_key_hints, + "full_path": None, + "exists": False, + "is_urn": True, + "connected": is_active, + } + ) + continue + + # Existing logic for local filenames + resolved = try_resolve_model_path( + value_str, + categories_to_try_for_widget, + available_models=available_models, + ) + + if resolved: + category, full_path = resolved + exists = os.path.exists(full_path) + else: + category = effective_category_hint or "unknown" + full_path = None + exists = False + + auto_download_capable = bool( + input_choice_matches_value + and input_choice_source in {"static", "hybrid", "workflow_schema"} + and input_choice_looks_like_model + ) + auto_download_candidate = bool(not exists and auto_download_capable) + + model_refs.append( + { + "node_id": node_id, + "node_type": node_type, + "widget_index": idx, + "widget_name": widget_name, + "original_path": value_str, + "node_title": node_title, + "category": category, + "category_hints": categories_to_try_for_widget or ([category] if category else []), + "folder_key_hints": model_widget_folder_key_hints + or categories_to_try_for_widget + or ([category] if category else []), + "full_path": full_path, + "exists": exists, + "input_choice_source": input_choice_source, + "input_choice_matches_value": input_choice_matches_value, + "auto_download_capable": auto_download_capable, + "auto_download_candidate": auto_download_candidate, + "is_urn": False, + "connected": is_active, + } + ) + + return model_refs diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py index da68e30..397c12f 100644 --- a/core/workflow_analyzer.py +++ b/core/workflow_analyzer.py @@ -7,7 +7,6 @@ Extracts model references from workflow JSON and identifies missing models. import hashlib import json import os -import re import threading import time from collections import OrderedDict @@ -32,694 +31,17 @@ _WORKFLOW_MODEL_INVENTORY_IGNORED_NODE_KEYS = { } -# Import folder_paths lazily - it may not be available until ComfyUI is initialized -try: - import folder_paths -except ImportError: - folder_paths = None - log.warning("Model Resolver: folder_paths not available yet - will retry later") - - -from .custom_nodes import ( - analyze_custom_node_references, - get_custom_node_model_adapter, -) -from .type_utils import ( - MODEL_EXTENSIONS, - URN_TYPE_MAP, - normalize_download_category, -) +from .workflow import references from .workflow.widgets import ( - NESTED_MODEL_KEYS, - NODE_TYPE_TO_CATEGORY_HINTS, _has_widget_input, - _ordered_unique_categories, _proxy_widget_name, _proxy_widget_node_id, _widget_item_name_candidates, - get_node_model_widget_category_hint, - get_node_output_category_hint, - get_widget_category_hint, get_widget_name_candidates, get_widget_name_hint, - is_workflow_model_widget_candidate, normalize_widget_name, ) -URN_REGEX = re.compile(r"^urn:air:([^:]+):([^:]+):([^:]+):(\d+)@(\d+)$") -from .workflow import dynamic_widgets - - -def is_static_or_hybrid_widget_choice(value: Any, choice_info: Dict[str, Any]) -> bool: - source = str(choice_info.get("source") or "").strip().lower() - if source not in {"static", "hybrid"}: - return False - - target = dynamic_widgets._normalize_choice_for_match(value) - if not target: - return False - - choices = choice_info.get("choices") or [] - return any( - dynamic_widgets._normalize_choice_for_match(choice) == target - for choice in choices - ) - - -def static_or_hybrid_choice_looks_like_model( - value: Any, choice_info: Dict[str, Any] -) -> bool: - if is_model_filename(value): - return True - - if not isinstance(value, str): - return False - - value_text = value.strip() - if not value_text: - return False - - if "/" in value_text or "\\" in value_text: - return True - - target = dynamic_widgets._normalize_choice_for_match(value_text) - choices = choice_info.get("choices") or [] - return any( - dynamic_widgets._normalize_choice_for_match(choice) == target - and is_model_filename(choice) - for choice in choices - ) - - -MODEL_WIDGET_PLACEHOLDERS = { - "", - ".none", - ".use_ckpt_clip", - ".use_ckpt_vae", - "none", - "[none]", - "null", - "undefined", - "default", - "auto", - "baked vae", - "included", - "(use same)", - "select the lora to add to the text", - "esam", - "pixel_space", - "taesd", - "taesdxl", - "taesd3", - "taef1", - "taef2", -} - - -def get_model_widget_category_hint( - node: Dict[str, Any], widget_index: int -) -> Optional[str]: - category_hints = get_model_widget_category_hints(node, widget_index) - return category_hints[0] if category_hints else None - - -def get_model_widget_category_hints( - node: Dict[str, Any], widget_index: int -) -> List[str]: - node_type = node.get("type", "") - indexed_category_hint = get_node_model_widget_category_hint(node_type, widget_index) - widget_category_hint = get_widget_category_hint(node, widget_index) - dynamic_category_hints = dynamic_widgets.get_dynamic_widget_category_hints( - node, widget_index - ) - - output_category_hint = get_node_output_category_hint(node) - output_widget_category_hint = ( - output_category_hint - if ( - not dynamic_category_hints - and (indexed_category_hint or widget_category_hint) - ) - else None - ) - - hints_list = [] - for h in dynamic_category_hints: - if h: - hints_list.append(normalize_download_category(h)) - - if dynamic_category_hints: - return _ordered_unique_categories(hints_list) - - if indexed_category_hint: - hints_list.append(indexed_category_hint) - - # If the widget category hint is just "diffusion_models", and we have a more specific output hint, - # let the output hint take priority. - if widget_category_hint == "diffusion_models" and output_widget_category_hint and output_widget_category_hint != "diffusion_models": - hints_list.append(output_widget_category_hint) - hints_list.append(widget_category_hint) - else: - if widget_category_hint: - hints_list.append(widget_category_hint) - if output_widget_category_hint: - hints_list.append(output_widget_category_hint) - - return _ordered_unique_categories(hints_list) - - -def get_effective_model_category_hint( - node: Dict[str, Any], widget_index: int -) -> Optional[str]: - return get_model_widget_category_hint( - node, widget_index - ) or NODE_TYPE_TO_CATEGORY_HINTS.get(node.get("type", "")) - - -def is_placeholder_model_value(value: Any) -> bool: - if not isinstance(value, str): - return False - - return value.strip().lower() in MODEL_WIDGET_PLACEHOLDERS - - -def should_scan_as_model_reference(value: Any, declared_model_widget: bool) -> bool: - """Detect references only when the value or its widget provides model evidence.""" - if isinstance(value, str) and URN_REGEX.match(value.strip()): - return True - - if not declared_model_widget or not isinstance(value, str): - return False - - if is_model_filename(value): - return True - - return bool(value.strip()) and not is_placeholder_model_value(value) - - -def is_model_filename(value: Any) -> bool: - """ - Check if a value looks like a model filename or URN. - - Args: - value: The value to check - - Returns: - True if it looks like a model filename or URN - """ - if not isinstance(value, str): - return False - - # Check model extension - _, ext = os.path.splitext(value.lower()) - if ext in MODEL_EXTENSIONS: - return True - - # Check URN format - return bool(URN_REGEX.match(value.strip())) - - -def _normalize_model_path_for_lookup(value: str) -> str: - """Normalize separators for comparison while preserving letter case.""" - if not isinstance(value, str): - return "" - - normalized = os.path.normpath(value.strip()) - if normalized == ".": - return "" - - return normalized.replace("\\", "/").strip("/") - - -def _resolve_from_available_models( - filename: str, - categories: Optional[List[str]], - available_models: Optional[List[Dict[str, Any]]], -) -> Optional[tuple[str, str]]: - """ - Resolve using scanner data with case-sensitive relative path matching. - - Windows accepts case-mismatched paths in os.path.exists(), but ComfyUI workflow - values need to match the actual folder/file casing to be safely reusable. - """ - if not available_models: - return None - - requested_key = _normalize_model_path_for_lookup(filename) - if not requested_key: - return None - - requested_is_absolute = os.path.isabs(filename) - if requested_is_absolute: - requested_key = _normalize_model_path_for_lookup(os.path.abspath(filename)) - - if categories is None: - category_order = [] - seen_categories = set() - for model in available_models: - category = model.get("category") - if category and category not in seen_categories: - seen_categories.add(category) - category_order.append(category) - else: - category_order = categories - - if not category_order: - return None - - for category in category_order: - normalized_category = normalize_download_category(category) - for model in available_models: - model_category = model.get("category") - if ( - model_category != category - and normalize_download_category(model_category) - != normalized_category - ): - continue - - if requested_is_absolute: - model_path = model.get("path") or "" - else: - model_path = model.get("relative_path") or model.get("filename") or "" - - if _normalize_model_path_for_lookup(model_path) != requested_key: - continue - - full_path = model.get("path") - if full_path and os.path.exists(full_path): - return (category, full_path) - - return None - - -def _category_has_exact_filename(category: str, filename: str) -> Optional[bool]: - """ - Check ComfyUI's filename list for an exact-case relative path match. - - Returns None when the filename list is unavailable, allowing callers to fall - back to the older filesystem existence path. - """ - if os.path.isabs(filename): - return None - - requested_key = _normalize_model_path_for_lookup(filename) - if not requested_key: - return False - - try: - available_filenames = folder_paths.get_filename_list(category) or [] - except Exception: - return None - - return any( - _normalize_model_path_for_lookup(available_filename) == requested_key - for available_filename in available_filenames - if isinstance(available_filename, str) - ) - - -def try_resolve_model_path( - value: str, - categories: List[str] = None, - available_models: Optional[List[Dict[str, Any]]] = None, -) -> Optional[tuple[str, str]]: - """ - Try to resolve a model path using folder_paths. - - Args: - value: The model filename/path to resolve - categories: Optional list of categories to try (if None, tries all) - available_models: Optional scanner results used for exact-case matching - - Returns: - Tuple of (category, full_path) if found, None otherwise - """ - if not isinstance(value, str) or not value.strip(): - return None - - # Remove any path separators that might indicate an absolute path prefix - # Workflows should store relative paths, but handle both cases - filename = value.strip() - - if categories is not None: - skip_categories = {"custom_nodes", "configs"} - categories = [c for c in categories if c not in skip_categories] - - resolved = _resolve_from_available_models(filename, categories, available_models) - if resolved: - return resolved - - if available_models is not None: - return None - - # Ensure folder_paths is available - global folder_paths - if folder_paths is None: - try: - import folder_paths as fp - - folder_paths = fp - except ImportError: - log.error("Model Resolver: folder_paths not available") - return None - - # If categories not provided, try all categories - if categories is None: - categories = list(folder_paths.folder_names_and_paths.keys()) - - # Skip non-model categories - skip_categories = {"custom_nodes", "configs"} - categories = [c for c in categories if c not in skip_categories] - - for category in categories: - try: - exact_filename = _category_has_exact_filename(category, filename) - if exact_filename is False: - continue - - full_path = folder_paths.get_full_path(category, filename) - if full_path and os.path.exists(full_path): - return (category, full_path) - except Exception: - continue - - return None - - -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. - - This scans widget values backed by model metadata or a model-like widget name. - Explicit model URNs are also accepted without widget metadata. - - Args: - node: Node dictionary from workflow JSON - - Returns: - List of model reference dictionaries: - { - 'node_id': node id, - 'node_type': node type, - 'widget_index': index in widgets_values, - 'original_path': original path from workflow, - 'category': model category (if found), - 'exists': True if model exists, - 'connected': True if node has any connected inputs/outputs - } - """ - model_refs = [] - node_id = node.get("id") - node_type = node.get("type", "") - node_title = str(node.get("title", "") or "").strip() - widgets_values = node.get("widgets_values", []) - - # Check if node is connected (has any inputs or outputs with links) - inputs = node.get("inputs", []) - outputs = node.get("outputs", []) - is_connected = any(inp.get("link") is not None for inp in inputs) or any( - out.get("links") and len(out.get("links", [])) > 0 for out in outputs - ) - - # Check if node is in bypass mode (mode 4) - node_mode = node.get("mode", 0) - is_bypassed = node_mode == 4 - - # Node is active if connected AND not bypassed - is_active = is_connected and not is_bypassed - - if not widgets_values: - return model_refs - - custom_model_refs = analyze_custom_node_references( - node, - available_models, - is_active=is_active, - get_widget_name_hint=get_widget_name_hint, - ) - if custom_model_refs is not None: - return custom_model_refs - - # For each widget value, check if it looks like a model file or URN - for idx, value in enumerate(widgets_values): - if dynamic_widgets.is_dynamic_non_model_widget(node, idx): - continue - widget_name = ( - dynamic_widgets.get_dynamic_serialized_widget_name(node, idx) - or get_widget_name_hint(node, idx) - ) - dynamic_category_hints = dynamic_widgets.get_dynamic_widget_category_hints( - node, idx - ) - if dynamic_category_hints and all( - normalize_download_category(category) - in dynamic_widgets.NON_MODEL_REFERENCE_CATEGORIES - for category in dynamic_category_hints - ): - continue - model_widget_folder_key_hints = dynamic_category_hints - model_widget_choice_info = dynamic_widgets.get_dynamic_widget_choice_info( - node, idx - ) - model_widget_category_hints = get_model_widget_category_hints(node, idx) - model_widget_category_hint = ( - model_widget_category_hints[0] if model_widget_category_hints else None - ) - input_choice_source = str( - model_widget_choice_info.get("source") or "unknown" - ).strip().lower() - static_input_choice_matches_value = is_static_or_hybrid_widget_choice( - value, - model_widget_choice_info, - ) - static_input_choice_looks_like_model = static_or_hybrid_choice_looks_like_model( - value, - model_widget_choice_info, - ) - if ( - input_choice_source in {"static", "hybrid"} - and static_input_choice_matches_value - and not static_input_choice_looks_like_model - ): - continue - output_category_hint = get_node_output_category_hint(node) - workflow_schema_model_candidate = bool( - input_choice_source == "unknown" - and output_category_hint - and is_workflow_model_widget_candidate(node, idx) - and is_model_filename(value) - ) - if workflow_schema_model_candidate: - input_choice_source = "workflow_schema" - input_choice_matches_value = bool( - static_input_choice_matches_value or workflow_schema_model_candidate - ) - input_choice_looks_like_model = bool( - static_input_choice_looks_like_model or workflow_schema_model_candidate - ) - schema_output_category_hint = ( - output_category_hint - if input_choice_matches_value - and input_choice_source in {"static", "hybrid", "workflow_schema"} - and input_choice_looks_like_model - else None - ) - effective_category_hint = ( - model_widget_category_hint - or NODE_TYPE_TO_CATEGORY_HINTS.get(node_type) - or schema_output_category_hint - ) - categories_to_try_for_widget = ( - model_widget_category_hints - if model_widget_category_hints - else ([effective_category_hint] if effective_category_hint else None) - ) - # A generic widget name such as ``model_name`` is useful for choosing a - # category, but it is not proof that an arbitrary STRING value is a local - # model reference. Some metadata/search nodes use that name for human- - # readable titles. Only folder-backed schemas, known widget indexes, or a - # model-typed combo/output pairing are strong enough to admit an - # extensionless value. - category_backed_model_widget = bool( - model_widget_folder_key_hints - or get_node_model_widget_category_hint(node_type, idx) - or schema_output_category_hint - ) - named_model_file_widget = bool( - is_workflow_model_widget_candidate(node, idx) - and is_model_filename(value) - ) - - if not should_scan_as_model_reference( - value, - declared_model_widget=( - category_backed_model_widget or named_model_file_widget - ), - ): - # Check for dict-type widget values containing model references. - if isinstance(value, dict): - for nested_key, nested_category_hint in NESTED_MODEL_KEYS.items(): - nested_value = value.get(nested_key) - if ( - not nested_value - or not isinstance(nested_value, str) - or not should_scan_as_model_reference( - nested_value, declared_model_widget=True - ) - ): - continue - - value_str = nested_value.strip() - nested_categories = ( - [nested_category_hint] if nested_category_hint else None - ) - - resolved = try_resolve_model_path( - value_str, - nested_categories, - available_models=available_models, - ) - if resolved: - category, full_path = resolved - exists = os.path.exists(full_path) - else: - category = nested_category_hint or "unknown" - full_path = None - exists = False - - ref = { - "node_id": node_id, - "node_type": node_type, - "widget_index": idx, - "widget_name": widget_name, - "original_path": value_str, - "node_title": node_title, - "category": category, - "category_hints": nested_categories or ([category] if category else []), - "folder_key_hints": nested_categories or ([category] if category else []), - "full_path": full_path, - "exists": exists, - "is_urn": False, - "connected": is_active, - "nested_key": nested_key, # Track nested key for updates - } - custom_adapter = get_custom_node_model_adapter(node) - if custom_adapter: - ref["custom_node_adapter"] = custom_adapter.adapter_id - - if nested_key == "lora": - strength = value.get("strength") - if strength is not None: - try: - ref["strength"] = float(strength) - except (TypeError, ValueError): - pass - - if isinstance(value.get("on"), bool): - ref["active"] = value.get("on") - - model_refs.append(ref) - continue - - value_str = str(value).strip() - - # Check if URN - urn_match = URN_REGEX.match(value_str) - if urn_match: - base, typ, provider, model_id, version_id = urn_match.groups() - category = ( - effective_category_hint - or URN_TYPE_MAP.get(typ.lower(), "unknown") - ) - urn_category_hints = ( - categories_to_try_for_widget - if categories_to_try_for_widget - else ([category] if category else []) - ) - urn_folder_key_hints = ( - model_widget_folder_key_hints - if model_widget_folder_key_hints - else urn_category_hints - ) - - model_refs.append( - { - "node_id": node_id, - "node_type": node_type, - "widget_index": idx, - "widget_name": widget_name, - "original_path": value_str, - "urn": { - "full": value_str, - "base": base, - "type": typ, - "provider": provider, - "model_id": int(model_id), - "version_id": int(version_id), - }, - "node_title": node_title, - "category": category, - "category_hints": urn_category_hints, - "folder_key_hints": urn_folder_key_hints, - "full_path": None, - "exists": False, - "is_urn": True, - "connected": is_active, - } - ) - continue - - # Existing logic for local filenames - resolved = try_resolve_model_path( - value_str, - categories_to_try_for_widget, - available_models=available_models, - ) - - if resolved: - category, full_path = resolved - exists = os.path.exists(full_path) - else: - category = effective_category_hint or "unknown" - full_path = None - exists = False - - auto_download_capable = bool( - input_choice_matches_value - and input_choice_source in {"static", "hybrid", "workflow_schema"} - and input_choice_looks_like_model - ) - auto_download_candidate = bool(not exists and auto_download_capable) - - model_refs.append( - { - "node_id": node_id, - "node_type": node_type, - "widget_index": idx, - "widget_name": widget_name, - "original_path": value_str, - "node_title": node_title, - "category": category, - "category_hints": categories_to_try_for_widget or ([category] if category else []), - "folder_key_hints": model_widget_folder_key_hints - or categories_to_try_for_widget - or ([category] if category else []), - "full_path": full_path, - "exists": exists, - "input_choice_source": input_choice_source, - "input_choice_matches_value": input_choice_matches_value, - "auto_download_capable": auto_download_capable, - "auto_download_candidate": auto_download_candidate, - "is_urn": False, - "connected": is_active, - } - ) - - return model_refs - def _get_node_by_id(nodes: List[Dict[str, Any]], node_id: Any) -> Optional[Dict[str, Any]]: node_id_text = str(node_id) @@ -789,7 +111,7 @@ def _promoted_widget_target_info( "node_title": str(node.get("title", "") or "").strip(), "widget_index": widget_index, "widget_name": widget_name, - "category": get_effective_model_category_hint(node, widget_index), + "category": references.get_effective_model_category_hint(node, widget_index), } @@ -937,7 +259,7 @@ def _apply_instance_promoted_widget_context( if not category or not original_path: return - resolved = try_resolve_model_path( + resolved = references.try_resolve_model_path( original_path, [category], available_models=available_models, @@ -1299,7 +621,7 @@ def analyze_workflow_models( model_refs = cached_node.get("refs", []) reused_nodes += 1 else: - model_refs = get_node_model_info( + model_refs = references.get_node_model_info( node, available_models=available_models, ) @@ -1379,7 +701,7 @@ def analyze_workflow_models( model_refs = cached_node.get("refs", []) reused_nodes += 1 else: - model_refs = get_node_model_info( + model_refs = references.get_node_model_info( node, available_models=available_models, ) diff --git a/tests/test_workflow_analyzer.py b/tests/test_workflow_analyzer.py index b5d1880..ec3f76a 100644 --- a/tests/test_workflow_analyzer.py +++ b/tests/test_workflow_analyzer.py @@ -7,14 +7,14 @@ from unittest.mock import patch from core import resolver as resolver_core from core import workflow_analyzer from core.scanner import invalidate_model_files_cache, scan_directory -from core.workflow import dynamic_widgets +from core.workflow import dynamic_widgets, references +from core.workflow.dynamic_widgets import get_lora_model_strength from core.workflow_analyzer import ( analyze_workflow_models, get_workflow_model_inventory, identify_missing_models, invalidate_workflow_model_inventory_cache, ) -from core.workflow.dynamic_widgets import get_lora_model_strength def _workflow_with_model(model_path): @@ -992,7 +992,7 @@ class WorkflowCategoryHintsTests(unittest.TestCase): self.assertEqual("checkpoints", refs[0]["category"]) def test_node_type_to_category_hints_is_populated(self): - from core.workflow_analyzer import NODE_TYPE_TO_CATEGORY_HINTS + from core.workflow.widgets import NODE_TYPE_TO_CATEGORY_HINTS # Verify standard loader mappings are correctly generated self.assertEqual("checkpoints", NODE_TYPE_TO_CATEGORY_HINTS.get("CheckpointLoaderSimple")) self.assertEqual("checkpoints", NODE_TYPE_TO_CATEGORY_HINTS.get("CheckpointLoader")) @@ -1417,9 +1417,9 @@ class WorkflowModelInventoryCacheTests(unittest.TestCase): return_value=[], ), patch.object( - workflow_analyzer, + references, "get_node_model_info", - wraps=workflow_analyzer.get_node_model_info, + wraps=references.get_node_model_info, ) as get_node_info, ): get_workflow_model_inventory(first_workflow) @@ -1465,9 +1465,9 @@ class WorkflowModelInventoryCacheTests(unittest.TestCase): return_value=[], ), patch.object( - workflow_analyzer, + references, "get_node_model_info", - wraps=workflow_analyzer.get_node_model_info, + wraps=references.get_node_model_info, ) as get_node_info, ): get_workflow_model_inventory(first_workflow) @@ -1501,9 +1501,9 @@ class WorkflowModelInventoryCacheTests(unittest.TestCase): return_value=[], ), patch.object( - workflow_analyzer, + references, "get_node_model_info", - wraps=workflow_analyzer.get_node_model_info, + wraps=references.get_node_model_info, ) as get_node_info, ): get_workflow_model_inventory(first_workflow) @@ -1564,9 +1564,9 @@ class WorkflowModelInventoryCacheTests(unittest.TestCase): return_value=[], ), patch.object( - workflow_analyzer, + references, "get_node_model_info", - wraps=workflow_analyzer.get_node_model_info, + wraps=references.get_node_model_info, ) as get_node_info, ): get_workflow_model_inventory(