Split off function from execute and simplified tests

This commit is contained in:
Deluxe233 2026-02-14 11:33:17 -05:00
parent 6172f7facc
commit 2581ac6376
3 changed files with 70 additions and 94 deletions

View File

@ -411,6 +411,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)
@ -451,32 +480,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 = []
for is_subgraph, result in cached_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))
resolved_outputs = resolve_subgraph_outputs(cached_results, unique_id, output_is_list, execution_list)
output_data = merge_result_data(resolved_outputs, class_def)
output_ui = []
del pending_subgraph_results[unique_id]

View File

@ -524,55 +524,35 @@ class TestExecution:
# 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", value1=0.1)
list_out = g.node("TestListExpansionResult")
output = g.node("SaveImage", images=list_out.out(0))
output_constant = g.node("SaveImage", images=list_out.out(1))
# return list of one image (list of one link)
result = client.run(g)
images = result.get_images(output)
assert len(images) == 1, "Should have 1 image"
assert numpy.array(images[0]).min() == 25 and numpy.array(images[0]).max() == 25, "First image should be 0.1"
images_constant = result.get_images(output_constant)
assert len(images_constant) == 1, "Should have 1 image"
assert numpy.array(images_constant[0]).min() == 255 and numpy.array(images_constant[0]).max() == 255, "Image should be white"
# test return list of two images (list of two links)
list_out.set_input("value2", 0.2)
result = client.run(g)
images = result.get_images(output)
assert len(images) == 2, "Should have 2 images"
assert numpy.array(images[0]).min() == 25 and numpy.array(images[0]).max() == 25, "First image should be 0.1"
assert numpy.array(images[1]).min() == 51 and numpy.array(images[1]).max() == 51, "Second image should be 0.2"
images_constant = result.get_images(output_constant)
assert len(images_constant) == 1, "Should have 1 image"
assert numpy.array(images_constant[0]).min() == 255 and numpy.array(images_constant[0]).max() == 255, "Image should be white"
# test mixed links and non-link values in returned list
list_out.set_input("value3", 0.3)
result = client.run(g)
images = result.get_images(output)
assert len(images) == 3, "Should have 3 images"
assert numpy.array(images[0]).min() == 25 and numpy.array(images[0]).max() == 25, "First image should be 0.1"
assert numpy.array(images[1]).min() == 51 and numpy.array(images[1]).max() == 51, "Second image should be 0.2"
assert numpy.array(images[2]).min() == 76 and numpy.array(images[2]).max() == 76, "Third image should be 0.3"
images_constant = result.get_images(output_constant)
assert len(images_constant) == 1, "Should have 1 image"
assert numpy.array(images_constant[0]).min() == 255 and numpy.array(images_constant[0]).max() == 255, "Image should be white"
# test returning list of a single link from an list output subnode
list_out.set_input("value4", 0.4)
result = client.run(g)
images = result.get_images(output)
assert len(images) == 4, "Should have 4 images"
assert numpy.array(images[0]).min() == 25 and numpy.array(images[0]).max() == 25, "First image should be 0.1"
assert numpy.array(images[1]).min() == 51 and numpy.array(images[1]).max() == 51, "Second image should be 0.2"
assert numpy.array(images[2]).min() == 76 and numpy.array(images[2]).max() == 76, "Third image should be 0.3"
assert numpy.array(images[3]).min() == 102 and numpy.array(images[3]).max() == 102, "Fourth image should be 0.4"
images_constant = result.get_images(output_constant)
assert len(images_constant) == 1, "Should have 1 image"
assert numpy.array(images_constant[0]).min() == 255 and numpy.array(images_constant[0]).max() == 255, "Image should be white"
# 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

View File

@ -359,35 +359,15 @@ class TestListExpansionResult:
CATEGORY = "Testing/Nodes"
def result_as_list(self, **kwargs):
g = GraphBuilder()
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) == 1:
image1 = g.node("StubConstantImage", value=values[0], height=512, width=512, batch_size=1)
return {
"result": ([image1.out(0)], white.out(0)),
"expand": g.finalize(),
}
elif len(values) == 2:
image1 = g.node("StubConstantImage", value=values[0], height=512, width=512, batch_size=1)
image2 = g.node("StubConstantImage", value=values[1], height=512, width=512, batch_size=1)
return {
"result": ([image1.out(0), image2.out(0)], white.out(0)),
"expand": g.finalize(),
}
elif len(values) == 3:
image1 = g.node("StubConstantImage", value=values[0], height=512, width=512, batch_size=1)
image2 = g.node("StubConstantImage", value=values[1], height=512, width=512, batch_size=1)
image3 = torch.ones(1, 512, 512, 3) * values[2]
return {
"result": ([image1.out(0), image2.out(0), image3], white.out(0)),
"expand": g.finalize(),
}
elif len(values) == 4:
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)
@ -396,6 +376,18 @@ class TestListExpansionResult:
"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