Merge 2581ac6376 into 9642e4407b

execution.py

@@ -410,6 +410,35 @@ def format_value(x):
     else:
         return str(x)
 
+def resolve_subgraph_outputs(subgraph_results, unique_id, output_is_list, execution_list):
+    resolved_outputs = []
+    for is_subgraph, result in subgraph_results:
+        if not is_subgraph:
+            resolved_outputs.append(result)
+        else:
+            resolved_output = []
+            for i, _result in enumerate(result):
+                if not output_is_list[i]:
+                    if is_link(_result):
+                        source_node, source_output = _result[0], _result[1]
+                        node_cached = execution_list.get_cache(source_node, unique_id)
+                        if node_cached.outputs[source_output]:
+                            resolved_output.append(node_cached.outputs[source_output][0])
+                    else:
+                        resolved_output.append(_result)
+                else:
+                    _resolved = []
+                    for output in _result:
+                        if is_link(output):
+                            source_node, source_output = output[0], output[1]
+                            node_cached = execution_list.get_cache(source_node, unique_id)
+                            _resolved.extend(node_cached.outputs[source_output])
+                        else:
+                            _resolved.append(output)
+                    resolved_output.append(_resolved)
+            resolved_outputs.append(tuple(resolved_output))
+    return resolved_outputs
+
 async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs):
     unique_id = current_item
     real_node_id = dynprompt.get_real_node_id(unique_id)
@@ -418,6 +447,9 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
     inputs = dynprompt.get_node(unique_id)['inputs']
     class_type = dynprompt.get_node(unique_id)['class_type']
     class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
+    output_is_list = [False] * len(class_def.RETURN_TYPES)
+    if hasattr(class_def, "OUTPUT_IS_LIST"):
+        output_is_list = class_def.OUTPUT_IS_LIST
     cached = caches.outputs.get(unique_id)
     if cached is not None:
         if server.client_id is not None:
@@ -447,22 +479,7 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
             output_data, output_ui, has_subgraph = get_output_from_returns(results, class_def)
         elif unique_id in pending_subgraph_results:
             cached_results = pending_subgraph_results[unique_id]
-            resolved_outputs = []
+            resolved_outputs = resolve_subgraph_outputs(cached_results, unique_id, output_is_list, execution_list)
-            for is_subgraph, result in cached_results:
-                if not is_subgraph:
-                    resolved_outputs.append(result)
-                else:
-                    resolved_output = []
-                    for r in result:
-                        if is_link(r):
-                            source_node, source_output = r[0], r[1]
-                            node_cached = execution_list.get_cache(source_node, unique_id)
-                            for o in node_cached.outputs[source_output]:
-                                resolved_output.append(o)
-
-                        else:
-                            resolved_output.append(r)
-                    resolved_outputs.append(tuple(resolved_output))
             output_data = merge_result_data(resolved_outputs, class_def)
             output_ui = []
             del pending_subgraph_results[unique_id]
@@ -570,9 +587,14 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
                         if hasattr(class_def, 'OUTPUT_NODE') and class_def.OUTPUT_NODE == True:
                             new_output_ids.append(node_id)
                     for i in range(len(node_outputs)):
-                        if is_link(node_outputs[i]):
+                        # Consider a returned list if output_is_list on the parent node
-                            from_node_id, from_socket = node_outputs[i][0], node_outputs[i][1]
+                        _node_outputs = node_outputs[i]
-                            new_output_links.append((from_node_id, from_socket))
+                        if not output_is_list[i]:
+                            _node_outputs = [_node_outputs]
+                        for node_output in _node_outputs:
+                            if is_link(node_output):
+                                from_node_id, from_socket = node_output[0], node_output[1]
+                                new_output_links.append((from_node_id, from_socket))
                     cached_outputs.append((True, node_outputs))
             new_node_ids = set(new_node_ids)
             for cache in caches.all:

tests/execution/test_execution.py

@@ -522,6 +522,38 @@ class TestExecution:
         for i in range(3):
             assert numpy.array(images_literal[i]).min() == 255 and numpy.array(images_literal[i]).max() == 255, "All images should be white"
 
+    # Tests functionality of defining OUTPUT_IS_LIST for expanding nodes.
+    def test_output_is_list_expansion_results(self, client: ComfyClient, builder: GraphBuilder):
+        def assert_image_values(images):
+            if len(images) >= 1:
+                assert numpy.array(images[0]).min() == 25 and numpy.array(images[0]).max() == 25, "First image should be 0.1"
+            if len(images) >= 2:
+                assert numpy.array(images[1]).min() == 51 and numpy.array(images[1]).max() == 51, "Second image should be 0.2"
+            if len(images) >= 3:
+                assert numpy.array(images[2]).min() == 76 and numpy.array(images[2]).max() == 76, "Third image should be 0.3"
+            if len(images) >= 4:
+                assert numpy.array(images[3]).min() == 102 and numpy.array(images[3]).max() == 102, "Fourth image should be 0.4"
+
+        def assert_constant_image(images):
+            assert len(images) == 1, "Should have 1 image"
+            assert numpy.array(images[0]).min() == 255 and numpy.array(images[0]).max() == 255, "Image should be white"
+
+        g = builder
+        list_out = g.node("TestListExpansionResult")
+        output = g.node("SaveImage", images=list_out.out(0))
+        output_constant = g.node("SaveImage", images=list_out.out(1))
+
+        # Run and check results for each new value added as input
+        for i in range(1, 5):
+            # value1 = 0.1, value2 = 0.2, etc.
+            list_out.set_input(f"value{i}", 0.1 * i)
+            result = client.run(g)
+            images = result.get_images(output)
+            assert len(images) == i, f"Should have {i} image(s)"
+            assert_image_values(images)
+            images_constant = result.get_images(output_constant)
+            assert_constant_image(images_constant)
+
     def test_mixed_lazy_results(self, client: ComfyClient, builder: GraphBuilder):
         g = builder
         val_list = g.node("TestMakeListNode", value1=0.0, value2=0.5, value3=1.0)

tests/execution/testing_nodes/testing-pack/specific_tests.py

@@ -338,6 +338,57 @@ class TestMixedExpansionReturns:
                 "expand": g.finalize(),
             }
 
+class TestListExpansionResult:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "value1": ("FLOAT",),
+            },
+            "optional": {
+                "value2": ("FLOAT",),
+                "value3": ("FLOAT",),
+                "value4": ("FLOAT",),
+            },
+        }
+    
+    RETURN_TYPES = ("IMAGE", "IMAGE")
+    FUNCTION = "result_as_list"
+    OUTPUT_IS_LIST = (True, False)
+    
+    CATEGORY = "Testing/Nodes"
+
+    def result_as_list(self, **kwargs):
+        values = []
+        for i in range(4):
+            key = f"value{i+1}"
+            if key in kwargs:
+                values.append(kwargs[key])
+        g = GraphBuilder()
+        white = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
+
+        if len(values) >= 4:
+            list_out = g.node("TestMakeListNode")
+            for i, value in enumerate(values):
+                image = g.node("StubConstantImage", value=value, height=512, width=512, batch_size=1)
+                list_out.set_input(f"value{i+1}", image.out(0))
+            return {
+                "result": ([list_out.out(0)], white.out(0)),
+                "expand": g.finalize(),
+            }
+        
+        images_out = []
+        if len(values) >= 1:
+            images_out.append(g.node("StubConstantImage", value=values[0], height=512, width=512, batch_size=1).out(0))
+        if len(values) >= 2:
+            images_out.append(g.node("StubConstantImage", value=values[1], height=512, width=512, batch_size=1).out(0))
+        if len(values) >= 3:
+            images_out.append(torch.ones(1, 512, 512, 3) * values[2])
+        return {
+            "result": (images_out, white.out(0)),
+            "expand": g.finalize(),
+        }
+
 class TestSamplingInExpansion:
     @classmethod
     def INPUT_TYPES(cls):
@@ -494,6 +545,7 @@ TEST_NODE_CLASS_MAPPINGS = {
     "TestCustomValidation5": TestCustomValidation5,
     "TestDynamicDependencyCycle": TestDynamicDependencyCycle,
     "TestMixedExpansionReturns": TestMixedExpansionReturns,
+    "TestListExpansionResult": TestListExpansionResult,
     "TestSamplingInExpansion": TestSamplingInExpansion,
     "TestSleep": TestSleep,
     "TestParallelSleep": TestParallelSleep,
@@ -512,6 +564,7 @@ TEST_NODE_DISPLAY_NAME_MAPPINGS = {
     "TestCustomValidation5": "Custom Validation 5",
     "TestDynamicDependencyCycle": "Dynamic Dependency Cycle",
     "TestMixedExpansionReturns": "Mixed Expansion Returns",
+    "TestListExpansionResult": "Output is List Expansion Result",
     "TestSamplingInExpansion": "Sampling In Expansion",
     "TestSleep": "Test Sleep",
     "TestParallelSleep": "Test Parallel Sleep",