Files

846 lines
26 KiB
Python

"""Dynamic workflow widget schemas and category inference."""
import re
import threading
from typing import Any, Dict, List, Optional, Set
from ..log_system import create_module_logger
from ..type_utils import unique_ordered_strings
from .widgets import (
_ordered_unique_categories,
get_widget_name_candidates,
normalize_widget_name,
)
try:
import folder_paths
except ImportError:
folder_paths = None
log = create_module_logger(__name__)
_DYNAMIC_CATEGORY_SENTINEL_PREFIX = "__model_resolver_folder_category__"
_DYNAMIC_NODE_WIDGET_CATEGORY_CACHE: Dict[str, Dict[str, Any]] = {}
_DYNAMIC_NODE_WIDGET_CATEGORY_LOCK = threading.RLock()
NON_MODEL_REFERENCE_CATEGORIES = {"custom_nodes", "configs"}
# These ComfyUI INPUT_TYPES entries become widgets in widgets_values. Typed graph
# inputs like MODEL or CLIP are links, so they should not shift widget indexes.
_WIDGET_INPUT_TYPES = {"BOOLEAN", "COMBO", "FLOAT", "INT", "STRING"}
def _merge_category_hints(
target: Dict[Any, List[str]], key: Any, categories: List[str]
) -> None:
if not categories:
return
target[key] = _ordered_unique_categories(target.get(key, []) + categories)
def _merge_choice_info_values(
current: Dict[str, Any],
incoming: Optional[Dict[str, Any]],
) -> Dict[str, Any]:
if not isinstance(incoming, dict):
return current
sources = _ordered_unique_categories(
[
source
for source in [current.get("source"), incoming.get("source")]
if str(source or "").lower() != "unknown"
]
)
if "hybrid" in sources or ("folder_paths" in sources and "static" in sources):
source = "hybrid"
else:
source = sources[0] if sources else "unknown"
return {
"source": source,
"choices": _ordered_unique_categories(
list(current.get("choices") or []) + list(incoming.get("choices") or [])
),
}
def _merge_choice_info(
target: Dict[Any, Dict[str, Any]], key: Any, info: Dict[str, Any]
) -> None:
if not isinstance(info, dict):
return
target[key] = _merge_choice_info_values(target.get(key, {}), info)
def _summarize_choice_info_for_log(info_by_key: Dict[Any, Dict[str, Any]]) -> Dict[Any, Dict[str, Any]]:
summary: Dict[Any, Dict[str, Any]] = {}
if not isinstance(info_by_key, dict):
return summary
for key, info in info_by_key.items():
if not isinstance(info, dict):
continue
summary[key] = {
"source": info.get("source", "unknown"),
"choice_count": len(info.get("choices") or []),
}
return summary
def _summarize_dynamic_hints_for_log(hints: Dict[str, Any]) -> Dict[str, Any]:
return {
"widget_names": hints.get("widget_names", []),
"by_name": hints.get("by_name", {}),
"by_index": hints.get("by_index", {}),
"serialized_name_by_index": hints.get("serialized_name_by_index", {}),
"non_model_by_index": hints.get("non_model_by_index", []),
"has_generated_widgets": hints.get("has_generated_widgets", False),
"choice_info_by_name": _summarize_choice_info_for_log(
hints.get("choice_info_by_name", {})
),
"choice_info_by_index": _summarize_choice_info_for_log(
hints.get("choice_info_by_index", {})
),
}
def _get_folder_paths_module() -> Any:
global folder_paths
if folder_paths is not None:
return folder_paths
try:
import folder_paths as fp
folder_paths = fp
return folder_paths
except Exception:
return None
def _get_comfy_node_class(node_type: str) -> Any:
if not node_type:
return None
try:
import nodes as comfy_nodes
except Exception:
return None
node_class_mappings = getattr(comfy_nodes, "NODE_CLASS_MAPPINGS", {}) or {}
if not isinstance(node_class_mappings, dict):
return None
return node_class_mappings.get(node_type)
def _extract_categories_from_value(
value: Any,
sentinel_to_category: Dict[str, str],
depth: int = 0,
seen: Optional[Set[int]] = None,
) -> List[str]:
if depth > 8 or not sentinel_to_category:
return []
if isinstance(value, str):
return _ordered_unique_categories(
[
category
for sentinel, category in sentinel_to_category.items()
if category and sentinel in value
]
)
if not isinstance(value, (dict, list, tuple, set)):
return []
if seen is None:
seen = set()
value_id = id(value)
if value_id in seen:
return []
seen.add(value_id)
categories: List[str] = []
if isinstance(value, dict):
for key, nested_value in value.items():
categories.extend(
_extract_categories_from_value(
key, sentinel_to_category, depth + 1, seen
)
)
categories.extend(
_extract_categories_from_value(
nested_value, sentinel_to_category, depth + 1, seen
)
)
else:
for item in value:
categories.extend(
_extract_categories_from_value(
item, sentinel_to_category, depth + 1, seen
)
)
return _ordered_unique_categories(categories)
def _flatten_combo_choice_values(
value: Any,
depth: int = 0,
seen: Optional[Set[int]] = None,
) -> List[str]:
if depth > 6:
return []
if isinstance(value, str):
return [value]
if not isinstance(value, (list, tuple, set)):
return []
if seen is None:
seen = set()
value_id = id(value)
if value_id in seen:
return []
seen.add(value_id)
choices: List[str] = []
for item in value:
choices.extend(_flatten_combo_choice_values(item, depth + 1, seen))
return _ordered_unique_categories(choices)
def _get_combo_choice_values(spec: Any) -> List[str]:
if not isinstance(spec, (list, tuple)) or not spec:
return []
input_type = spec[0]
if not isinstance(input_type, (list, tuple, set)):
return []
return _flatten_combo_choice_values(input_type)
def _choice_contains_category_sentinel(
value: Any,
sentinel_to_category: Dict[str, str],
) -> bool:
text = str(value or "")
return any(sentinel and sentinel in text for sentinel in sentinel_to_category)
def _get_widget_choice_info_from_spec(
spec: Any,
sentinel_to_category: Dict[str, str],
) -> Dict[str, Any]:
choices = _get_combo_choice_values(spec)
if not choices:
return {}
has_folder_choices = any(
_choice_contains_category_sentinel(choice, sentinel_to_category)
for choice in choices
)
static_choices = [
choice
for choice in choices
if not _choice_contains_category_sentinel(choice, sentinel_to_category)
]
if has_folder_choices and static_choices:
source = "hybrid"
elif has_folder_choices:
source = "folder_paths"
else:
source = "static"
return {
"source": source,
"choices": static_choices,
}
def _is_input_type_widget_spec(spec: Any) -> bool:
if not isinstance(spec, (list, tuple)) or not spec:
return False
input_type = spec[0]
if isinstance(input_type, (list, tuple)):
return True
if isinstance(input_type, str):
return input_type.strip().upper() in _WIDGET_INPUT_TYPES
return False
def _input_type_generates_control_after_widget(spec: Any) -> bool:
if not isinstance(spec, (list, tuple)) or len(spec) < 2:
return False
options = spec[1]
return isinstance(options, dict) and options.get("control_after_generate") is True
def _iter_widget_input_type_entries(
input_types: Dict[str, Any],
) -> List[tuple[str, Any]]:
entries: List[tuple[str, Any]] = []
if not isinstance(input_types, dict):
return entries
for section_name in ("required", "optional"):
section = input_types.get(section_name, {})
if not isinstance(section, dict):
continue
entries.extend(section.items())
return entries
def _empty_dynamic_hints() -> Dict[str, Any]:
return {
"widget_names": [],
"by_name": {},
"by_index": {},
"serialized_name_by_index": {},
"non_model_by_index": [],
"has_generated_widgets": False,
"choice_info_by_name": {},
"choice_info_by_index": {},
}
def _schema_input_name(input_obj: Any) -> str:
for key in ("id", "name", "display_name"):
value = getattr(input_obj, key, None)
if value:
return str(value).strip()
return ""
def _schema_input_io_type(input_obj: Any) -> str:
get_io_type = getattr(input_obj, "get_io_type", None)
if callable(get_io_type):
try:
return str(get_io_type() or "").strip().upper()
except Exception:
pass
return str(getattr(input_obj, "io_type", "") or "").strip().upper()
def _is_schema_input_widget(input_obj: Any) -> bool:
io_type = _schema_input_io_type(input_obj)
return io_type in _WIDGET_INPUT_TYPES
def _iter_schema_input_entries(schema: Any) -> List[tuple[str, Any, bool]]:
inputs = getattr(schema, "inputs", None)
if not isinstance(inputs, list):
return []
entries: List[tuple[str, Any, bool]] = []
for input_obj in inputs:
expanded_inputs = []
get_all = getattr(input_obj, "get_all", None)
if callable(get_all):
try:
expanded_inputs = get_all()
except Exception:
expanded_inputs = []
if not expanded_inputs:
expanded_inputs = [input_obj]
for expanded_input in expanded_inputs:
entries.append(
(
_schema_input_name(expanded_input),
expanded_input,
_is_schema_input_widget(expanded_input),
)
)
return entries
def _schema_input_choice_info(
input_obj: Any,
sentinel_to_category: Dict[str, str],
) -> Dict[str, Any]:
options = getattr(input_obj, "options", None)
if not isinstance(options, (list, tuple, set)):
return {}
return _get_widget_choice_info_from_spec((list(options),), sentinel_to_category)
def _merge_dynamic_hint_entry(
hints: Dict[str, Any],
normalized_name: str,
widget_index: Optional[int],
categories: List[str],
choice_info: Dict[str, Any],
) -> None:
if normalized_name:
hints["widget_names"] = _ordered_unique_categories(
[*list(hints.get("widget_names") or []), normalized_name]
)
_merge_category_hints(hints["by_name"], normalized_name, categories)
if choice_info:
_merge_choice_info(
hints["choice_info_by_name"],
normalized_name,
choice_info,
)
if widget_index is not None:
if categories:
_merge_category_hints(hints["by_index"], widget_index, categories)
if choice_info:
_merge_choice_info(
hints["choice_info_by_index"],
widget_index,
choice_info,
)
def _build_dynamic_node_widget_category_hints(node_type: str) -> Dict[str, Any]:
empty_hints = _empty_dynamic_hints()
node_class = _get_comfy_node_class(node_type)
if node_class is None:
return empty_hints
input_types_getter = getattr(node_class, "INPUT_TYPES", None)
schema_getter = getattr(node_class, "define_schema", None)
if not callable(input_types_getter) and not callable(schema_getter):
return empty_hints
folder_paths_module = _get_folder_paths_module()
get_filename_list = getattr(folder_paths_module, "get_filename_list", None)
sentinel_to_category: Dict[str, str] = {}
def traced_get_filename_list(category: Any, *args: Any, **kwargs: Any) -> List[str]:
category_name = str(category or "").strip()
sentinel = (
f"{_DYNAMIC_CATEGORY_SENTINEL_PREFIX}"
f"{len(sentinel_to_category)}.safetensors"
)
sentinel_to_category[sentinel] = category_name
return [sentinel]
input_types_func = getattr(input_types_getter, "__func__", input_types_getter)
schema_func = getattr(schema_getter, "__func__", schema_getter)
patched_globals: List[tuple[Dict[str, Any], str, Any]] = []
patched_folder_paths = False
def patch_get_filename_list_global(callable_obj: Any) -> None:
if not callable(callable_obj) or not callable(get_filename_list):
return
globals_dict = getattr(callable_obj, "__globals__", {})
if not isinstance(globals_dict, dict):
return
for global_name, global_value in list(globals_dict.items()):
if global_value is get_filename_list:
globals_dict[global_name] = traced_get_filename_list
patched_globals.append((globals_dict, global_name, global_value))
try:
if folder_paths_module is not None and callable(get_filename_list):
folder_paths_module.get_filename_list = traced_get_filename_list
patched_folder_paths = True
patch_get_filename_list_global(input_types_func)
patch_get_filename_list_global(schema_func)
input_types = input_types_getter() if callable(input_types_getter) else None
schema = schema_getter() if callable(schema_getter) else None
except Exception as exc:
log.debug(
f"Could not infer dynamic model widget categories for {node_type}: {exc}"
)
return empty_hints
finally:
for global_scope, global_name, global_value in patched_globals:
global_scope[global_name] = global_value
if patched_folder_paths:
folder_paths_module.get_filename_list = get_filename_list
if not isinstance(input_types, dict) and schema is None:
return empty_hints
hints = _empty_dynamic_hints()
widget_index = 0
for input_name, spec in _iter_widget_input_type_entries(input_types):
categories = _extract_categories_from_value(spec, sentinel_to_category)
choice_info = _get_widget_choice_info_from_spec(spec, sentinel_to_category)
normalized_name = normalize_widget_name(input_name)
is_widget = _is_input_type_widget_spec(spec)
entry_widget_index = widget_index if is_widget else None
if is_widget:
hints["serialized_name_by_index"][widget_index] = normalized_name
_merge_dynamic_hint_entry(
hints,
normalized_name,
entry_widget_index,
categories,
choice_info,
)
if is_widget:
widget_index += 1
if _input_type_generates_control_after_widget(spec):
hints["has_generated_widgets"] = True
hints["serialized_name_by_index"][
widget_index
] = "control_after_generate"
hints["non_model_by_index"].append(widget_index)
widget_index += 1
for input_name, input_obj, is_widget in _iter_schema_input_entries(schema):
categories = _extract_categories_from_value(
getattr(input_obj, "options", None),
sentinel_to_category,
)
choice_info = _schema_input_choice_info(input_obj, sentinel_to_category)
normalized_name = normalize_widget_name(input_name)
_merge_dynamic_hint_entry(
hints,
normalized_name,
widget_index if is_widget else None,
categories,
choice_info,
)
if is_widget:
widget_index += 1
if (
hints["by_name"]
or hints["by_index"]
or hints["choice_info_by_name"]
or hints["choice_info_by_index"]
):
log.debug(
"Inferred dynamic model widget categories for "
f"{node_type}: {_summarize_dynamic_hints_for_log(hints)}"
)
return hints
def get_dynamic_node_widget_category_hints(node_type: str) -> Dict[str, Any]:
if not node_type:
return _empty_dynamic_hints()
with _DYNAMIC_NODE_WIDGET_CATEGORY_LOCK:
cached = _DYNAMIC_NODE_WIDGET_CATEGORY_CACHE.get(node_type)
if cached is not None:
return cached
hints = _build_dynamic_node_widget_category_hints(node_type)
# Only cache after ComfyUI has exposed the node class and folder_paths. If
# analysis runs unusually early, a later call can still infer the hints.
if _get_comfy_node_class(node_type) is not None and _get_folder_paths_module():
_DYNAMIC_NODE_WIDGET_CATEGORY_CACHE[node_type] = hints
return hints
CONTROL_AFTER_GENERATE_VALUES = {
"fixed",
"increment",
"decrement",
"randomize",
}
def _normalized_index_mapping(value: Any) -> Dict[int, str]:
if not isinstance(value, dict):
return {}
normalized: Dict[int, str] = {}
for raw_index, raw_name in value.items():
try:
index = int(raw_index)
except (TypeError, ValueError):
continue
if index < 0:
continue
normalized[index] = str(raw_name or "").strip()
return normalized
def _dynamic_serialized_widget_layout(
node: Dict[str, Any], hints: Dict[str, Any]
) -> tuple[Dict[int, str], Set[int]]:
names_by_index = _normalized_index_mapping(
hints.get("serialized_name_by_index", {})
)
non_model_indices = {
int(index)
for index in hints.get("non_model_by_index", [])
if str(index).isdigit()
}
if hints.get("has_generated_widgets"):
return names_by_index, non_model_indices
widgets_values = node.get("widgets_values", [])
if not isinstance(widgets_values, list) or not names_by_index:
return {}, set()
expected_widget_count = max(names_by_index) + 1
remaining_extra_widgets = len(widgets_values) - expected_widget_count
if remaining_extra_widgets <= 0:
return {}, set()
effective_names: Dict[int, str] = {}
effective_non_model_indices: Set[int] = set()
inserted_widget_count = 0
for base_index in sorted(names_by_index):
widget_name = names_by_index[base_index]
actual_index = base_index + inserted_widget_count
effective_names[actual_index] = widget_name
normalized_name = normalize_widget_name(widget_name)
control_index = actual_index + 1
seed_value = (
widgets_values[actual_index]
if actual_index < len(widgets_values)
else None
)
control_value = (
str(widgets_values[control_index] or "").strip().lower()
if control_index < len(widgets_values)
else ""
)
is_seed_widget = normalized_name == "seed" or normalized_name.endswith(
"_seed"
)
if (
remaining_extra_widgets > 0
and is_seed_widget
and isinstance(seed_value, int)
and not isinstance(seed_value, bool)
and control_value in CONTROL_AFTER_GENERATE_VALUES
):
effective_names[control_index] = "control_after_generate"
effective_non_model_indices.add(control_index)
inserted_widget_count += 1
remaining_extra_widgets -= 1
if not effective_non_model_indices:
return {}, set()
return effective_names, effective_non_model_indices
def get_dynamic_serialized_widget_name(
node: Dict[str, Any], widget_index: int
) -> str:
hints = get_dynamic_node_widget_category_hints(str(node.get("type", "") or ""))
names_by_index, _non_model_indices = _dynamic_serialized_widget_layout(
node, hints
)
return names_by_index.get(widget_index, "")
LORA_MODEL_STRENGTH_WIDGET_NAMES = (
"strength_model",
"lora_model_strength",
"model_strength",
"lora_strength",
"strength",
)
LEGACY_LORA_STRENGTH_NEXT_WIDGET_TYPES = {
"LoraLoader",
"LoraLoaderModelOnly",
"LoraLoaderBypass",
"LoraLoaderBypassModelOnly",
"CreateHookLora",
"CreateHookLoraModelOnly",
}
def _coerce_lora_strength(value: Any) -> Optional[float]:
if isinstance(value, bool):
return None
try:
return float(value)
except (TypeError, ValueError):
return None
def get_lora_model_strength(
node: Dict[str, Any], lora_widget_index: Optional[int] = None
) -> Optional[float]:
"""Return the model strength associated with a LoRA widget in a workflow node."""
widgets_values = node.get("widgets_values", [])
if not isinstance(widgets_values, list):
return None
hints = get_dynamic_node_widget_category_hints(str(node.get("type", "") or ""))
schema_names = _normalized_index_mapping(
hints.get("serialized_name_by_index", {})
)
effective_names, _non_model_indices = _dynamic_serialized_widget_layout(
node, hints
)
if effective_names:
schema_names = effective_names
values_by_name: Dict[str, Any] = {}
names_by_index: Dict[int, List[str]] = {}
for widget_index, value in enumerate(widgets_values):
names = []
schema_name = schema_names.get(widget_index)
if schema_name:
names.append(schema_name)
names.extend(get_widget_name_candidates(node, widget_index))
normalized_names = unique_ordered_strings(
[normalize_widget_name(name) for name in names]
)
names_by_index[widget_index] = normalized_names
for normalized_name in normalized_names:
if normalized_name and normalized_name not in values_by_name:
values_by_name[normalized_name] = value
if isinstance(lora_widget_index, int):
for widget_name in names_by_index.get(lora_widget_index, []):
indexed_match = re.fullmatch(r"lora_(\d+)_name", widget_name)
if not indexed_match:
continue
lora_number = indexed_match.group(1)
mode = str(values_by_name.get("mode") or "").strip().lower()
if mode == "simple":
indexed_strength_names = [
f"lora_{lora_number}_strength",
f"lora_{lora_number}_model_strength",
]
else:
indexed_strength_names = [
f"lora_{lora_number}_model_strength",
f"lora_{lora_number}_strength",
]
for strength_name in indexed_strength_names:
strength = _coerce_lora_strength(
values_by_name.get(strength_name)
)
if strength is not None:
return strength
for strength_name in LORA_MODEL_STRENGTH_WIDGET_NAMES:
if strength_name not in values_by_name:
continue
strength = _coerce_lora_strength(values_by_name[strength_name])
if strength is not None:
return strength
node_type = str(node.get("type", "") or "")
if (
node_type in LEGACY_LORA_STRENGTH_NEXT_WIDGET_TYPES
and isinstance(lora_widget_index, int)
and 0 <= lora_widget_index + 1 < len(widgets_values)
):
return _coerce_lora_strength(widgets_values[lora_widget_index + 1])
return None
def is_dynamic_non_model_widget(node: Dict[str, Any], widget_index: int) -> bool:
hints = get_dynamic_node_widget_category_hints(str(node.get("type", "") or ""))
_names_by_index, non_model_indices = _dynamic_serialized_widget_layout(
node, hints
)
return widget_index in non_model_indices
def get_dynamic_widget_category_hints(
node: Dict[str, Any], widget_index: int
) -> List[str]:
hints = get_dynamic_node_widget_category_hints(str(node.get("type", "") or ""))
if is_dynamic_non_model_widget(node, widget_index):
return []
categories: List[str] = []
by_name = hints.get("by_name", {})
serialized_name = get_dynamic_serialized_widget_name(node, widget_index)
normalized_candidates = (
[normalize_widget_name(serialized_name)]
if serialized_name
else [
normalize_widget_name(candidate)
for candidate in get_widget_name_candidates(node, widget_index)
]
)
if isinstance(by_name, dict):
for candidate in normalized_candidates:
categories.extend(by_name.get(candidate, []))
known_widget_names = set(hints.get("widget_names") or [])
if any(candidate in known_widget_names for candidate in normalized_candidates):
return _ordered_unique_categories(categories)
by_index = hints.get("by_index", {})
if isinstance(by_index, dict):
categories.extend(by_index.get(widget_index, []))
return _ordered_unique_categories(categories)
def _normalize_choice_for_match(value: Any) -> str:
return (
str(value or "")
.strip()
.replace("\\", "/")
.lower()
)
def get_dynamic_widget_choice_info(
node: Dict[str, Any], widget_index: int
) -> Dict[str, Any]:
hints = get_dynamic_node_widget_category_hints(str(node.get("type", "") or ""))
info: Dict[str, Any] = {"source": "unknown", "choices": []}
if is_dynamic_non_model_widget(node, widget_index):
return info
by_name = hints.get("choice_info_by_name", {})
serialized_name = get_dynamic_serialized_widget_name(node, widget_index)
normalized_candidates = (
[normalize_widget_name(serialized_name)]
if serialized_name
else [
normalize_widget_name(candidate)
for candidate in get_widget_name_candidates(node, widget_index)
]
)
if isinstance(by_name, dict):
for candidate in normalized_candidates:
candidate_info = by_name.get(candidate, {})
if isinstance(candidate_info, dict):
info = _merge_choice_info_values(info, candidate_info)
known_widget_names = set(hints.get("widget_names") or [])
if any(candidate in known_widget_names for candidate in normalized_candidates):
return info
by_index = hints.get("choice_info_by_index", {})
if isinstance(by_index, dict):
candidate_info = by_index.get(widget_index, {})
if isinstance(candidate_info, dict):
info = _merge_choice_info_values(info, candidate_info)
return info