Add a test for mixed node expansion

This test ensures that a node that returns a combination of expanded
subgraphs and literal values functions correctly.
This commit is contained in:
Jacob Segal 2024-04-21 15:39:20 -07:00
parent 1f06588905
commit 2dda3f2827
3 changed files with 55 additions and 0 deletions

View File

@ -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

View File

@ -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"

View File

@ -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",
}