From ac9bf820c85678f43e9caa1c318da556327dfcb2 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Fri, 24 Jul 2026 15:10:00 +0200 Subject: [PATCH] Handle control_after_generate widget indexes in workflow analyzer --- core/workflow_analyzer.py | 98 +++++++++++++++-- tests/test_workflow_analyzer.py | 183 ++++++++++++++++++++++++++++++++ 2 files changed, 270 insertions(+), 11 deletions(-) diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py index 21d652d..3607fa0 100644 --- a/core/workflow_analyzer.py +++ b/core/workflow_analyzer.py @@ -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" diff --git a/tests/test_workflow_analyzer.py b/tests/test_workflow_analyzer.py index 7dfd3ae..6724d87 100644 --- a/tests/test_workflow_analyzer.py +++ b/tests/test_workflow_analyzer.py @@ -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"],