Handle control_after_generate widget indexes in workflow analyzer

This commit is contained in:
Dariusz L
2026-07-24 15:10:00 +02:00
parent 03c545e297
commit ac9bf820c8
2 changed files with 270 additions and 11 deletions
+87 -11
View File
@@ -403,6 +403,9 @@ def _summarize_dynamic_hints_for_log(hints: Dict[str, Any]) -> Dict[str, Any]:
"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", {})
),
@@ -587,6 +590,14 @@ def _is_input_type_widget_spec(spec: Any) -> bool:
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]]:
@@ -608,6 +619,9 @@ def _empty_dynamic_hints() -> Dict[str, Any]:
"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": {},
}
@@ -777,7 +791,10 @@ def _build_dynamic_node_widget_category_hints(node_type: str) -> Dict[str, Any]:
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)
entry_widget_index = widget_index if _is_input_type_widget_spec(spec) else None
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,
@@ -785,8 +802,15 @@ def _build_dynamic_node_widget_category_hints(node_type: str) -> Dict[str, Any]:
categories,
choice_info,
)
if _is_input_type_widget_spec(spec):
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(
@@ -838,17 +862,50 @@ def get_dynamic_node_widget_category_hints(node_type: str) -> Dict[str, Any]:
return hints
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 ""))
if not hints.get("has_generated_widgets"):
return ""
names_by_index = hints.get("serialized_name_by_index", {})
if not isinstance(names_by_index, dict):
return ""
return str(
names_by_index.get(widget_index, names_by_index.get(str(widget_index), "")) or ""
).strip()
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 ""))
if not hints.get("has_generated_widgets"):
return False
non_model_indices = hints.get("non_model_by_index", [])
return widget_index in non_model_indices or str(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", {})
normalized_candidates = [
normalize_widget_name(candidate)
for candidate in get_widget_name_candidates(node, widget_index)
]
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, []))
@@ -906,11 +963,19 @@ def get_dynamic_widget_choice_info(
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", {})
normalized_candidates = [
normalize_widget_name(candidate)
for candidate in get_widget_name_candidates(node, widget_index)
]
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, {})
@@ -1443,7 +1508,12 @@ def get_node_model_info(
# For each widget value, check if it looks like a model file or URN
for idx, value in enumerate(widgets_values):
widget_name = get_widget_name_hint(node, idx)
if is_dynamic_non_model_widget(node, idx):
continue
widget_name = (
get_dynamic_serialized_widget_name(node, idx)
or get_widget_name_hint(node, idx)
)
dynamic_category_hints = get_dynamic_widget_category_hints(node, idx)
if dynamic_category_hints and all(
normalize_download_category(category)
@@ -1468,6 +1538,12 @@ def get_node_model_info(
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"
+183
View File
@@ -190,6 +190,189 @@ class WorkflowAnalyzerCategoryHintTests(unittest.TestCase):
hints["by_name"]["model_name"],
)
def test_control_after_generate_keeps_serialized_widget_indexes_aligned(self):
fake_folder_paths = SimpleNamespace()
def get_filename_list(category):
return [f"{category}.safetensors"]
fake_folder_paths.get_filename_list = get_filename_list
class ControlledIndexLoader:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"csv_filename_path": ("STRING", {}),
"index": (
"INT",
{"default": 0, "control_after_generate": True},
),
"start_ckpt_name": (
fake_folder_paths.get_filename_list("checkpoints"),
),
"count": ("INT", {"default": 1}),
}
}
with (
patch.object(
workflow_analyzer,
"_get_comfy_node_class",
return_value=ControlledIndexLoader,
),
patch.object(
workflow_analyzer,
"_get_folder_paths_module",
return_value=fake_folder_paths,
),
):
hints = workflow_analyzer._build_dynamic_node_widget_category_hints(
"ControlledIndexLoader"
)
self.assertTrue(hints["has_generated_widgets"])
self.assertEqual(
{
0: "csv_filename_path",
1: "index",
2: "control_after_generate",
3: "start_ckpt_name",
4: "count",
},
hints["serialized_name_by_index"],
)
self.assertEqual([2], hints["non_model_by_index"])
self.assertNotIn(2, hints["by_index"])
self.assertEqual(["checkpoints"], hints["by_index"][3])
def test_control_after_generate_is_skipped_and_following_model_is_detected(self):
dynamic_hints = {
"widget_names": [
"csv_filename_path",
"index",
"start_ckpt_name",
"count",
],
"by_name": {"start_ckpt_name": ["checkpoints"]},
"by_index": {3: ["checkpoints"]},
"serialized_name_by_index": {
0: "csv_filename_path",
1: "index",
2: "control_after_generate",
3: "start_ckpt_name",
4: "count",
},
"non_model_by_index": [2],
"has_generated_widgets": True,
"choice_info_by_name": {
"start_ckpt_name": {
"source": "folder_paths",
"choices": [],
}
},
"choice_info_by_index": {
3: {"source": "folder_paths", "choices": []}
},
}
workflow = {
"nodes": [
{
"id": 503,
"type": "CSV Reader X Checkpoint",
"inputs": [
{
"name": "csv_filename_path",
"type": "STRING",
"widget": {"name": "csv_filename_path"},
},
{
"name": "index",
"type": "INT",
"widget": {"name": "index"},
},
{
"name": "start_ckpt_name",
"type": "COMBO",
"widget": {"name": "start_ckpt_name"},
},
{
"name": "count",
"type": "INT",
"widget": {"name": "count"},
},
],
"widgets_values": [
"",
0,
"randomize",
r"LTXV\model.safetensors",
1,
],
"outputs": [],
}
]
}
with patch(
"core.workflow_analyzer.get_dynamic_node_widget_category_hints",
return_value=dynamic_hints,
):
refs = analyze_workflow_models(workflow, available_models=[])
self.assertEqual(1, len(refs))
self.assertEqual(3, refs[0]["widget_index"])
self.assertEqual("start_ckpt_name", refs[0]["widget_name"])
self.assertEqual(r"LTXV\model.safetensors", refs[0]["original_path"])
self.assertEqual("checkpoints", refs[0]["category"])
def test_non_model_static_choices_in_hybrid_model_dropdowns_are_skipped(self):
cases = [
("DependenciesEdit", "ckpt_name", "checkpoints", "Original"),
("DependenciesEdit", "vae_name", "vae", "Baked VAE"),
("easy XYInputs: ModelMergeBlocks", "vae_use", "vae", "Use Model 1"),
("RecipeModelPicker", "vae", "vae", "(Default)"),
]
for node_type, widget_name, category, value in cases:
with self.subTest(node_type=node_type, value=value):
dynamic_hints = {
"widget_names": [widget_name],
"by_name": {widget_name: [category]},
"by_index": {0: [category]},
"serialized_name_by_index": {},
"non_model_by_index": [],
"has_generated_widgets": False,
"choice_info_by_name": {
widget_name: {
"source": "hybrid",
"choices": [value],
}
},
"choice_info_by_index": {
0: {"source": "hybrid", "choices": [value]}
},
}
workflow = {
"nodes": [
{
"id": 504,
"type": node_type,
"widgets": [{"name": widget_name}],
"widgets_values": [value],
"outputs": [],
}
]
}
with patch(
"core.workflow_analyzer.get_dynamic_node_widget_category_hints",
return_value=dynamic_hints,
):
refs = analyze_workflow_models(workflow, available_models=[])
self.assertEqual([], refs)
def test_res4lyf_placeholders_are_not_model_references(self):
category_by_index = {
0: ["checkpoints", "diffusion_models"],