655 lines
22 KiB
Python
655 lines
22 KiB
Python
"""Model reference matching, category inference, and path resolution."""
|
|
|
|
import os
|
|
import re
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from ..contracts import ModelReference, ResolvedModel
|
|
from ..custom_nodes import (
|
|
analyze_custom_node_references,
|
|
get_custom_node_model_adapter,
|
|
)
|
|
from ..log_system import create_module_logger
|
|
from ..metadata_model_utils import normalize_models
|
|
from ..path_utils import normalize_string_values
|
|
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[ResolvedModel]],
|
|
) -> 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))
|
|
|
|
normalized_models = normalize_models(available_models)
|
|
|
|
if categories is None:
|
|
category_order = []
|
|
seen_categories = set()
|
|
for model in normalized_models:
|
|
category = model.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 normalized_models:
|
|
model_category = model.category
|
|
if (
|
|
model_category != category
|
|
and normalize_download_category(model_category) != normalized_category
|
|
):
|
|
continue
|
|
|
|
if requested_is_absolute:
|
|
model_path = model.path or ""
|
|
else:
|
|
model_path = model.relative_path or model.filename or ""
|
|
|
|
if _normalize_model_path_for_lookup(model_path) != requested_key:
|
|
continue
|
|
|
|
full_path = model.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 = normalize_string_values(
|
|
folder_paths.get_filename_list(category)
|
|
)
|
|
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[ResolvedModel]] = 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[ResolvedModel]] = None
|
|
) -> List[ModelReference]:
|
|
"""
|
|
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:
|
|
Typed model references. Provider/custom-node-specific fields are kept
|
|
in ``ModelReference.extra``.
|
|
"""
|
|
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 []
|
|
|
|
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 [
|
|
ref if isinstance(ref, ModelReference) else ModelReference.from_mapping(ref)
|
|
for ref in 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 [
|
|
ref if isinstance(ref, ModelReference) else ModelReference.from_mapping(ref)
|
|
for ref in model_refs
|
|
]
|