From 6e7818bee6de1fd3b97bf10783e59703936a9c7a Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Fri, 5 Jun 2026 17:12:25 +1000 Subject: [PATCH] fix: preserve delegate-only master config inputs (#86) Fixes #85 --- api/orchestration/prompt_transform.py | 206 +++++++++++ tests/test_prompt_transform.py | 495 ++++++++++++++++++++++++++ 2 files changed, 701 insertions(+) diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index 244bd6c..ebaa29f 100644 --- a/api/orchestration/prompt_transform.py +++ b/api/orchestration/prompt_transform.py @@ -125,6 +125,209 @@ def _find_upstream_nodes(prompt_obj, start_ids): return connected +_DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES = { + "PrimitiveBoolean", + "PrimitiveFloat", + "PrimitiveInt", + "PrimitiveNode", + "PrimitiveString", +} + +_DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES = { + "LoadImage", +} + +_DELEGATE_MASTER_SAFE_SCALAR_TYPES = {"BOOLEAN", "FLOAT", "INT", "STRING"} +_DELEGATE_MASTER_SAFE_LIST_TYPES = {"LIST"} + +# ComfyUI 0.23 exposes CreateList via the newer schema API rather than the +# legacy RETURN_TYPES/INPUT_TYPES attributes. Treat it as a safe config utility +# only after its connected inputs recursively prove safe. +_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES = {"CreateList"} +_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES = { + "CreateList": ("inputs.",), +} + +# Test hook. At runtime this stays None and the ComfyUI node registry is loaded lazily. +_DELEGATE_MASTER_NODE_CLASS_MAPPINGS = None + + +def _get_delegate_master_node_class_mappings(): + """Return ComfyUI node-class mappings when available.""" + if _DELEGATE_MASTER_NODE_CLASS_MAPPINGS is not None: + return _DELEGATE_MASTER_NODE_CLASS_MAPPINGS + try: + import nodes as comfy_nodes # type: ignore + except Exception: # pragma: no cover - depends on ComfyUI runtime imports + return {} + return getattr(comfy_nodes, "NODE_CLASS_MAPPINGS", {}) or {} + + +def _get_delegate_master_node_class(class_type): + mappings = _get_delegate_master_node_class_mappings() + return mappings.get(class_type) if isinstance(mappings, dict) else None + + +def _normalize_delegate_master_return_type(return_type): + if return_type is None: + return "" + return str(return_type).strip().upper() + + +def _delegate_master_type_is_safe_scalar(type_name): + return type_name in _DELEGATE_MASTER_SAFE_SCALAR_TYPES + + +def _delegate_master_type_is_safe_config(type_name): + return _delegate_master_type_is_safe_scalar(type_name) or type_name in _DELEGATE_MASTER_SAFE_LIST_TYPES + + +def _delegate_master_output_is_safe_scalar(class_type, output_index): + """Return True when a registered node output is lightweight config data.""" + if class_type in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES: + return True + node_class = _get_delegate_master_node_class(class_type) + return_types = getattr(node_class, "RETURN_TYPES", ()) if node_class is not None else () + try: + output_type = return_types[int(output_index)] + except (IndexError, TypeError, ValueError): + return False + return _delegate_master_type_is_safe_config(_normalize_delegate_master_return_type(output_type)) + + +def _get_delegate_master_input_types(class_type): + node_class = _get_delegate_master_node_class(class_type) + input_types = getattr(node_class, "INPUT_TYPES", None) if node_class is not None else None + if callable(input_types): + try: + input_types = input_types() + except TypeError: + return {} + return input_types if isinstance(input_types, dict) else {} + + +def _normalize_delegate_master_input_type(input_spec): + if isinstance(input_spec, (list, tuple)) and input_spec: + return _normalize_delegate_master_return_type(input_spec[0]) + return _normalize_delegate_master_return_type(input_spec) + + +def _delegate_master_input_is_safe_scalar(class_type, input_name): + """Return True when a registered downstream input expects config data.""" + for prefix in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES.get(class_type, ()): + if input_name.startswith(prefix): + return True + input_types = _get_delegate_master_input_types(class_type) + for section_name in ("required", "optional"): + section = input_types.get(section_name, {}) + if isinstance(section, dict) and input_name in section: + input_type = _normalize_delegate_master_input_type(section[input_name]) + return _delegate_master_type_is_safe_config(input_type) + return False + + +def _is_delegate_master_always_retained_upstream_node(node): + if not isinstance(node, dict): + return False + class_type = node.get("class_type") + return isinstance(class_type, str) and class_type in _DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES + + +def _is_delegate_master_retained_upstream_node(node, output_index=0): + """Return True for lightweight upstream nodes safe to keep on the master.""" + if not isinstance(node, dict): + return False + class_type = node.get("class_type") + if not isinstance(class_type, str): + return False + return ( + class_type in _DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES + or class_type.startswith("Primitive") + or _delegate_master_output_is_safe_scalar(class_type, output_index) + ) + + +def _collect_delegate_master_retained_upstream_branch( + prompt_obj, + node_id, + output_index, + memo, + visiting, +): + """Return safe retained branch nodes, or None when the branch is not safe.""" + node_id = str(node_id) + cache_key = (node_id, output_index) + if cache_key in memo: + cached = memo[cache_key] + return None if cached is None else set(cached) + if cache_key in visiting: + memo[cache_key] = None + return None + + node = prompt_obj.get(node_id) + if not _is_delegate_master_retained_upstream_node(node, output_index): + memo[cache_key] = None + return None + + visiting.add(cache_key) + retained = {node_id} + inputs = node.get("inputs", {}) if isinstance(node, dict) else {} + class_type = node.get("class_type") if isinstance(node, dict) else None + for input_name, value in inputs.items(): + if not (isinstance(value, list) and len(value) == 2): + continue + if not _delegate_master_input_is_safe_scalar(class_type, input_name): + visiting.remove(cache_key) + memo[cache_key] = None + return None + source_id = str(value[0]) + branch = _collect_delegate_master_retained_upstream_branch( + prompt_obj, + source_id, + value[1], + memo, + visiting, + ) + if branch is None: + visiting.remove(cache_key) + memo[cache_key] = None + return None + retained.update(branch) + + visiting.remove(cache_key) + memo[cache_key] = frozenset(retained) + return retained + + +def _find_delegate_master_retained_upstream_nodes(prompt_obj, start_ids): + """Return lightweight upstream nodes needed by kept delegate-master nodes.""" + connected = set() + memo = {} + for node_id in start_ids: + node = prompt_obj.get(str(node_id)) or {} + inputs = node.get("inputs", {}) + class_type = node.get("class_type") if isinstance(node, dict) else None + for input_name, value in inputs.items(): + if not (isinstance(value, list) and len(value) == 2): + continue + source_node = prompt_obj.get(str(value[0])) + if _is_delegate_master_always_retained_upstream_node(source_node): + connected.add(str(value[0])) + continue + if not _delegate_master_input_is_safe_scalar(class_type, input_name): + continue + branch = _collect_delegate_master_retained_upstream_branch( + prompt_obj, + value[0], + value[1], + memo, + set(), + ) + if branch is not None: + connected.update(branch) + return connected + + def prune_prompt_for_worker(prompt_obj): """Prune worker prompt to distributed nodes and their upstream dependencies.""" collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector") @@ -167,6 +370,9 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids): downstream = _find_downstream_nodes(prompt_obj, collector_ids) nodes_to_keep = set(collector_ids) nodes_to_keep.update(downstream) + nodes_to_keep.update( + _find_delegate_master_retained_upstream_nodes(prompt_obj, nodes_to_keep) + ) pruned_prompt = {} for node_id in nodes_to_keep: diff --git a/tests/test_prompt_transform.py b/tests/test_prompt_transform.py index 0b5bf6f..2bd200a 100644 --- a/tests/test_prompt_transform.py +++ b/tests/test_prompt_transform.py @@ -296,6 +296,501 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase): empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"] self.assertEqual(len(empty_nodes), 2) + def test_preserves_primitive_string_for_downstream_required_input(self): + """Delegate-only master keeps primitive inputs needed by SaveImage.""" + save_image = type( + "SaveImage", + (), + { + "INPUT_TYPES": classmethod( + lambda cls: { + "required": { + "images": ("IMAGE",), + "filename_prefix": ("STRING",), + } + } + ) + }, + ) + previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None) + pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {"SaveImage": save_image} + try: + prompt = { + "8": {"class_type": "KSampler", "inputs": {}}, + "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}}, + "15": {"class_type": "PrimitiveString", "inputs": {"value": "test/input_bug"}}, + "9": { + "class_type": "SaveImage", + "inputs": { + "images": ["11", 0], + "filename_prefix": ["15", 0], + }, + }, + } + + result = pt.prepare_delegate_master_prompt(prompt, ["11"]) + finally: + pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings + + self.assertIn("15", result) + self.assertEqual(result["15"], prompt["15"]) + self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["15", 0]) + + def test_does_not_preserve_non_primitive_upstream_for_collector(self): + prompt = _delegate_prompt() + result = pt.prepare_delegate_master_prompt(prompt, ["3"]) + + self.assertNotIn("2", result) + self.assertNotEqual(result["3"]["inputs"]["images"], ["2", 0]) + + def test_preserves_load_image_for_switch_alternate_required_input(self): + """Delegate-only master keeps LoadImage inputs needed by switches.""" + prompt = { + "848": {"class_type": "LoadImage", "inputs": {"image": "input_bug2_00003_.png"}}, + "862": {"class_type": "VAEDecode", "inputs": {}}, + "854": {"class_type": "DistributedCollector", "inputs": {"images": ["862", 0]}}, + "850": { + "class_type": "ComfySwitchNode", + "inputs": { + "on_false": ["848", 0], + "on_true": ["854", 0], + }, + }, + "851": {"class_type": "PreviewImage", "inputs": {"images": ["850", 0]}}, + } + + result = pt.prepare_delegate_master_prompt(prompt, ["854"]) + + self.assertIn("848", result) + self.assertEqual(result["850"]["inputs"]["on_false"], ["848", 0]) + self.assertEqual(result["850"]["inputs"]["on_true"], ["854", 0]) + self.assertNotIn("862", result) + + def test_preserves_registered_string_utility_subgraph_for_downstream_required_input(self): + """Delegate-only master keeps scalar utility chains used by SaveImage.""" + string_concat = type( + "StringConcatenate", + (), + { + "RETURN_TYPES": ("STRING",), + "INPUT_TYPES": classmethod( + lambda cls: { + "required": { + "string_a": ("STRING",), + "string_b": ("STRING",), + } + } + ), + }, + ) + save_image = type( + "SaveImage", + (), + { + "INPUT_TYPES": classmethod( + lambda cls: { + "required": { + "images": ("IMAGE",), + "filename_prefix": ("STRING",), + } + } + ) + }, + ) + previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None) + pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = { + "SaveImage": save_image, + "StringConcatenate": string_concat, + } + try: + prompt = { + "8": {"class_type": "KSampler", "inputs": {}}, + "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}}, + "15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}}, + "16": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug3"}}, + "17": { + "class_type": "StringConcatenate", + "inputs": { + "string_a": ["15", 0], + "string_b": ["16", 0], + }, + }, + "9": { + "class_type": "SaveImage", + "inputs": { + "images": ["11", 0], + "filename_prefix": ["17", 0], + }, + }, + } + + result = pt.prepare_delegate_master_prompt(prompt, ["11"]) + finally: + pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings + + self.assertIn("15", result) + self.assertIn("16", result) + self.assertIn("17", result) + self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["17", 0]) + self.assertEqual(result["17"]["inputs"]["string_a"], ["15", 0]) + self.assertEqual(result["17"]["inputs"]["string_b"], ["16", 0]) + + def test_preserves_registered_multi_string_join_subgraph_for_downstream_required_input(self): + """Delegate-only master keeps multi-input scalar utility chains.""" + join_string_multi = type( + "JoinStringMulti", + (), + { + "RETURN_TYPES": ("STRING",), + "INPUT_TYPES": classmethod( + lambda cls: { + "required": {"string_1": ("STRING",)}, + "optional": { + "string_2": ("STRING",), + "string_3": ("STRING",), + }, + } + ), + }, + ) + save_image = type( + "SaveImage", + (), + { + "INPUT_TYPES": classmethod( + lambda cls: { + "required": { + "images": ("IMAGE",), + "filename_prefix": ("STRING",), + } + } + ) + }, + ) + previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None) + pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = { + "JoinStringMulti": join_string_multi, + "SaveImage": save_image, + } + try: + prompt = { + "8": {"class_type": "KSampler", "inputs": {}}, + "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}}, + "15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}}, + "16": {"class_type": "PrimitiveString", "inputs": {"value": "input"}}, + "17": {"class_type": "PrimitiveString", "inputs": {"value": "bug4"}}, + "18": { + "class_type": "JoinStringMulti", + "inputs": { + "string_1": ["15", 0], + "string_2": ["16", 0], + "string_3": ["17", 0], + }, + }, + "9": { + "class_type": "SaveImage", + "inputs": { + "images": ["11", 0], + "filename_prefix": ["18", 0], + }, + }, + } + + result = pt.prepare_delegate_master_prompt(prompt, ["11"]) + finally: + pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings + + self.assertIn("15", result) + self.assertIn("16", result) + self.assertIn("17", result) + self.assertIn("18", result) + self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["18", 0]) + + def test_does_not_preserve_scalar_utility_with_heavy_upstream_dependency(self): + """Scalar utility nodes are retained only when their full input branch is safe.""" + string_concat = type( + "StringConcatenate", + (), + { + "RETURN_TYPES": ("STRING",), + "INPUT_TYPES": classmethod( + lambda cls: { + "required": { + "string_a": ("STRING",), + "string_b": ("STRING",), + } + } + ), + }, + ) + save_image = type( + "SaveImage", + (), + { + "INPUT_TYPES": classmethod( + lambda cls: { + "required": { + "images": ("IMAGE",), + "filename_prefix": ("STRING",), + } + } + ) + }, + ) + previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None) + pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = { + "SaveImage": save_image, + "StringConcatenate": string_concat, + } + try: + prompt = { + "8": {"class_type": "KSampler", "inputs": {}}, + "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}}, + "15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}}, + "17": { + "class_type": "StringConcatenate", + "inputs": { + "string_a": ["15", 0], + "string_b": ["8", 0], + }, + }, + "9": { + "class_type": "SaveImage", + "inputs": { + "images": ["11", 0], + "filename_prefix": ["17", 0], + }, + }, + } + + result = pt.prepare_delegate_master_prompt(prompt, ["11"]) + finally: + pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings + + self.assertNotIn("8", result) + self.assertNotIn("17", result) + self.assertNotIn("filename_prefix", result["9"]["inputs"]) + + def test_does_not_preserve_scalar_output_for_non_scalar_downstream_input(self): + """Scalar outputs are retained only for scalar/config downstream inputs.""" + string_provider = type("StringProvider", (), {"RETURN_TYPES": ("STRING",)}) + image_consumer = type( + "ImageConsumer", + (), + { + "INPUT_TYPES": classmethod( + lambda cls: { + "required": { + "images": ("IMAGE",), + "mask": ("IMAGE",), + } + } + ) + }, + ) + previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None) + pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = { + "ImageConsumer": image_consumer, + "StringProvider": string_provider, + } + try: + prompt = { + "8": {"class_type": "KSampler", "inputs": {}}, + "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}}, + "17": {"class_type": "StringProvider", "inputs": {}}, + "9": { + "class_type": "ImageConsumer", + "inputs": { + "images": ["11", 0], + "mask": ["17", 0], + }, + }, + } + + result = pt.prepare_delegate_master_prompt(prompt, ["11"]) + finally: + pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings + + self.assertNotIn("17", result) + self.assertNotIn("mask", result["9"]["inputs"]) + + def test_preserves_scalar_list_join_subgraph_for_downstream_required_input(self): + """Delegate-only master keeps list-shaped scalar config chains.""" + create_list = type( + "CreateList", + (), + { + "RETURN_TYPES": ("LIST",), + "INPUT_TYPES": classmethod( + lambda cls: { + "required": {"inputs.input0": ("STRING",)}, + "optional": { + "inputs.input1": ("STRING",), + "inputs.input2": ("STRING",), + }, + } + ), + }, + ) + string_data_list_join = type( + "StringDataListJoin", + (), + { + "RETURN_TYPES": ("STRING",), + "INPUT_TYPES": classmethod( + lambda cls: { + "required": { + "strings": ("LIST",), + "delimiter": ("STRING",), + } + } + ), + }, + ) + save_image = type( + "SaveImage", + (), + { + "INPUT_TYPES": classmethod( + lambda cls: { + "required": { + "images": ("IMAGE",), + "filename_prefix": ("STRING",), + } + } + ) + }, + ) + previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None) + pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = { + "CreateList": create_list, + "SaveImage": save_image, + "StringDataListJoin": string_data_list_join, + } + try: + prompt = { + "8": {"class_type": "KSampler", "inputs": {}}, + "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}}, + "15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}}, + "17": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug"}}, + "29": {"class_type": "PrimitiveString", "inputs": {"value": "new"}}, + "28": { + "class_type": "CreateList", + "inputs": { + "inputs.input0": ["15", 0], + "inputs.input1": ["17", 0], + "inputs.input2": ["29", 0], + }, + }, + "32": { + "class_type": "StringDataListJoin", + "inputs": { + "strings": ["28", 0], + "delimiter": "/", + }, + }, + "9": { + "class_type": "SaveImage", + "inputs": { + "images": ["11", 0], + "filename_prefix": ["32", 0], + }, + }, + } + + result = pt.prepare_delegate_master_prompt(prompt, ["11"]) + finally: + pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings + + self.assertIn("15", result) + self.assertIn("17", result) + self.assertIn("28", result) + self.assertIn("29", result) + self.assertIn("32", result) + self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["32", 0]) + self.assertEqual(result["32"]["inputs"]["strings"], ["28", 0]) + self.assertEqual(result["28"]["inputs"]["inputs.input0"], ["15", 0]) + self.assertNotIn("8", result) + + def test_preserves_builtin_create_list_string_data_list_join_subgraph(self): + """Delegate-only master handles ComfyUI 0.23 CreateList data-list output.""" + string_data_list_join = type( + "StringDataListJoin", + (), + { + "RETURN_TYPES": ("STRING",), + "INPUT_IS_LIST": True, + "INPUT_TYPES": classmethod( + lambda cls: { + "required": { + "strings": ("STRING", {"forceInput": True}), + "sep": ("STRING", {"default": " "}), + } + } + ), + }, + ) + save_image = type( + "SaveImage", + (), + { + "INPUT_TYPES": classmethod( + lambda cls: { + "required": { + "images": ("IMAGE",), + "filename_prefix": ("STRING",), + } + } + ) + }, + ) + previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None) + pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = { + "Basic data handling: StringDataListJoin": string_data_list_join, + "SaveImage": save_image, + } + try: + prompt = { + "8": {"class_type": "KSampler", "inputs": {}}, + "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}}, + "15": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug"}}, + "17": {"class_type": "PrimitiveString", "inputs": {"value": "test"}}, + "29": {"class_type": "PrimitiveString", "inputs": {"value": "new"}}, + "28": { + "class_type": "CreateList", + "inputs": { + "inputs.input0": ["17", 0], + "inputs.input1": ["15", 0], + "inputs.input2": ["29", 0], + }, + }, + "32": { + "class_type": "Basic data handling: StringDataListJoin", + "inputs": { + "strings": ["28", 0], + "sep": "/", + }, + }, + "9": { + "class_type": "SaveImage", + "inputs": { + "images": ["11", 0], + "filename_prefix": ["32", 0], + }, + }, + } + + result = pt.prepare_delegate_master_prompt(prompt, ["11"]) + finally: + pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings + + for node_id in ("15", "17", "28", "29", "32"): + self.assertIn(node_id, result) + self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["32", 0]) + self.assertEqual(result["32"]["inputs"]["strings"], ["28", 0]) + self.assertEqual(result["28"]["inputs"]["inputs.input0"], ["17", 0]) + self.assertNotIn("8", result) + def test_result_is_independent_copy(self): prompt = _delegate_prompt() result = pt.prepare_delegate_master_prompt(prompt, ["3"])