Robert Wojciechowski
@@ -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:
|
||||
|
||||
@@ -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"])
|
||||
|
||||
Reference in New Issue
Block a user