diff --git a/execution.py b/execution.py index 2aa939c1b..9db69cbd2 100644 --- a/execution.py +++ b/execution.py @@ -224,6 +224,7 @@ def get_output_data(obj, input_data_all, execution_block_cb=None, pre_execute_cb if isinstance(r, ExecutionBlocker): r = tuple([r] * len(obj.RETURN_TYPES)) results.append(r) + subgraph_results.append((None, r)) if has_subgraph: output = subgraph_results diff --git a/tests/inference/test_execution.py b/tests/inference/test_execution.py index 40de99833..1fb58d5e0 100644 --- a/tests/inference/test_execution.py +++ b/tests/inference/test_execution.py @@ -399,3 +399,22 @@ class TestExecution: expected = 255 // (2 ** iterations) assert numpy.array(result_image).min() == expected and numpy.array(result_image).max() == expected, "Image should be grey" assert result.did_run(is_changed) + + def test_mixed_expansion_returns(self, client: ComfyClient, builder: GraphBuilder): + g = builder + val_list = g.node("TestMakeListNode", value1=0.1, value2=0.2, value3=0.3) + mixed = g.node("TestMixedExpansionReturns", input1=val_list.out(0)) + output_dynamic = g.node("SaveImage", images=mixed.out(0)) + output_literal = g.node("SaveImage", images=mixed.out(1)) + + result = client.run(g) + images_dynamic = result.get_images(output_dynamic) + assert len(images_dynamic) == 3, "Should have 2 images" + assert numpy.array(images_dynamic[0]).min() == 25 and numpy.array(images_dynamic[0]).max() == 25, "First image should be 0.1" + assert numpy.array(images_dynamic[1]).min() == 51 and numpy.array(images_dynamic[1]).max() == 51, "Second image should be 0.2" + assert numpy.array(images_dynamic[2]).min() == 76 and numpy.array(images_dynamic[2]).max() == 76, "Third image should be 0.3" + + images_literal = result.get_images(output_literal) + assert len(images_literal) == 3, "Should have 2 images" + 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" diff --git a/tests/inference/testing_nodes/testing-pack/specific_tests.py b/tests/inference/testing_nodes/testing-pack/specific_tests.py index 03dbc8b42..8e8ce32ce 100644 --- a/tests/inference/testing_nodes/testing-pack/specific_tests.py +++ b/tests/inference/testing_nodes/testing-pack/specific_tests.py @@ -250,6 +250,39 @@ class TestDynamicDependencyCycle: "expand": g.finalize(), } +class TestMixedExpansionReturns: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "input1": ("FLOAT",), + }, + } + + RETURN_TYPES = ("IMAGE","IMAGE") + FUNCTION = "mixed_expansion_returns" + + CATEGORY = "Testing/Nodes" + + def mixed_expansion_returns(self, input1): + white_image = torch.ones([1, 512, 512, 3]) + if input1 <= 0.1: + return (torch.ones([1, 512, 512, 3]) * 0.1, white_image) + elif input1 <= 0.2: + return { + "result": (torch.ones([1, 512, 512, 3]) * 0.2, white_image), + } + else: + g = GraphBuilder() + mask = g.node("StubMask", value=0.3, height=512, width=512, batch_size=1) + black = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1) + white = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1) + mix = g.node("TestLazyMixImages", image1=black.out(0), image2=white.out(0), mask=mask.out(0)) + return { + "result": (mix.out(0), white_image), + "expand": g.finalize(), + } + TEST_NODE_CLASS_MAPPINGS = { "TestLazyMixImages": TestLazyMixImages, "TestVariadicAverage": TestVariadicAverage, @@ -259,6 +292,7 @@ TEST_NODE_CLASS_MAPPINGS = { "TestCustomValidation3": TestCustomValidation3, "TestCustomValidation4": TestCustomValidation4, "TestDynamicDependencyCycle": TestDynamicDependencyCycle, + "TestMixedExpansionReturns": TestMixedExpansionReturns, } TEST_NODE_DISPLAY_NAME_MAPPINGS = { @@ -270,4 +304,5 @@ TEST_NODE_DISPLAY_NAME_MAPPINGS = { "TestCustomValidation3": "Custom Validation 3", "TestCustomValidation4": "Custom Validation 4", "TestDynamicDependencyCycle": "Dynamic Dependency Cycle", + "TestMixedExpansionReturns": "Mixed Expansion Returns", }