mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-09-02 07:07:03 +08:00
745 lines
35 KiB
Python
745 lines
35 KiB
Python
from io import BytesIO
|
|
import numpy
|
|
from PIL import Image
|
|
import pytest
|
|
from pytest import fixture
|
|
import time
|
|
import torch
|
|
from typing import Union, Dict
|
|
import json
|
|
import subprocess
|
|
import websocket #NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
|
|
import uuid
|
|
import urllib.request
|
|
import urllib.parse
|
|
import urllib.error
|
|
from comfy_execution.graph_utils import GraphBuilder, Node
|
|
|
|
class RunResult:
|
|
def __init__(self, prompt_id: str):
|
|
self.outputs: Dict[str,Dict] = {}
|
|
self.runs: Dict[str,bool] = {}
|
|
self.prompt_id: str = prompt_id
|
|
|
|
def get_output(self, node: Node):
|
|
return self.outputs.get(node.id, None)
|
|
|
|
def did_run(self, node: Node):
|
|
return self.runs.get(node.id, False)
|
|
|
|
def get_images(self, node: Node):
|
|
output = self.get_output(node)
|
|
if output is None:
|
|
return []
|
|
return output.get('image_objects', [])
|
|
|
|
def get_prompt_id(self):
|
|
return self.prompt_id
|
|
|
|
class ComfyClient:
|
|
def __init__(self):
|
|
self.test_name = ""
|
|
|
|
def connect(self,
|
|
listen:str = '127.0.0.1',
|
|
port:Union[str,int] = 8188,
|
|
client_id: str = str(uuid.uuid4())
|
|
):
|
|
self.client_id = client_id
|
|
self.server_address = f"{listen}:{port}"
|
|
ws = websocket.WebSocket()
|
|
ws.connect("ws://{}/ws?clientId={}".format(self.server_address, self.client_id))
|
|
self.ws = ws
|
|
|
|
def queue_prompt(self, prompt):
|
|
p = {"prompt": prompt, "client_id": self.client_id}
|
|
data = json.dumps(p).encode('utf-8')
|
|
req = urllib.request.Request("http://{}/prompt".format(self.server_address), data=data)
|
|
return json.loads(urllib.request.urlopen(req).read())
|
|
|
|
def get_image(self, filename, subfolder, folder_type):
|
|
data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
|
|
url_values = urllib.parse.urlencode(data)
|
|
with urllib.request.urlopen("http://{}/view?{}".format(self.server_address, url_values)) as response:
|
|
return response.read()
|
|
|
|
def get_history(self, prompt_id=None, max_items=None):
|
|
if prompt_id:
|
|
url = "http://{}/history/{}".format(self.server_address, prompt_id)
|
|
else:
|
|
url = "http://{}/history".format(self.server_address)
|
|
if max_items is not None:
|
|
url += "?max_items={}".format(max_items)
|
|
with urllib.request.urlopen(url) as response:
|
|
return json.loads(response.read())
|
|
|
|
def get_ordered_history(self, max_items=None, offset=None):
|
|
url = "http://{}/history_v2".format(self.server_address)
|
|
params = {}
|
|
if max_items is not None:
|
|
params['max_items'] = str(max_items)
|
|
if offset is not None:
|
|
params['offset'] = str(offset)
|
|
if params:
|
|
url += "?" + urllib.parse.urlencode(params)
|
|
with urllib.request.urlopen(url) as response:
|
|
return json.loads(response.read())
|
|
|
|
def get_history_v2_for_prompt(self, prompt_id):
|
|
url = "http://{}/history_v2/{}".format(self.server_address, prompt_id)
|
|
with urllib.request.urlopen(url) as response:
|
|
return json.loads(response.read())
|
|
|
|
def clear_history(self):
|
|
data = json.dumps({"clear": True}).encode('utf-8')
|
|
req = urllib.request.Request("http://{}/history".format(self.server_address), data=data)
|
|
req.add_header('Content-Type', 'application/json')
|
|
urllib.request.urlopen(req)
|
|
|
|
def set_test_name(self, name):
|
|
self.test_name = name
|
|
|
|
def run(self, graph):
|
|
prompt = graph.finalize()
|
|
for node in graph.nodes.values():
|
|
if node.class_type == 'SaveImage':
|
|
node.inputs['filename_prefix'] = self.test_name
|
|
|
|
prompt_id = self.queue_prompt(prompt)['prompt_id']
|
|
result = RunResult(prompt_id)
|
|
while True:
|
|
out = self.ws.recv()
|
|
if isinstance(out, str):
|
|
message = json.loads(out)
|
|
if message['type'] == 'executing':
|
|
data = message['data']
|
|
if data['prompt_id'] != prompt_id:
|
|
continue
|
|
if data['node'] is None:
|
|
break
|
|
result.runs[data['node']] = True
|
|
elif message['type'] == 'execution_error':
|
|
raise Exception(message['data'])
|
|
elif message['type'] == 'execution_cached':
|
|
pass # Probably want to store this off for testing
|
|
|
|
history = self.get_history(prompt_id)[prompt_id]
|
|
for node_id in history['outputs']:
|
|
node_output = history['outputs'][node_id]
|
|
result.outputs[node_id] = node_output
|
|
images_output = []
|
|
if 'images' in node_output:
|
|
for image in node_output['images']:
|
|
image_data = self.get_image(image['filename'], image['subfolder'], image['type'])
|
|
image_obj = Image.open(BytesIO(image_data))
|
|
images_output.append(image_obj)
|
|
node_output['image_objects'] = images_output
|
|
|
|
return result
|
|
|
|
#
|
|
# Loop through these variables
|
|
#
|
|
@pytest.mark.execution
|
|
class TestExecution:
|
|
#
|
|
# Initialize server and client
|
|
#
|
|
@fixture(scope="class", autouse=True, params=[
|
|
# (use_lru, lru_size)
|
|
(False, 0),
|
|
(True, 0),
|
|
(True, 100),
|
|
])
|
|
def _server(self, args_pytest, request):
|
|
# Start server
|
|
pargs = [
|
|
'python','main.py',
|
|
'--output-directory', args_pytest["output_dir"],
|
|
'--listen', args_pytest["listen"],
|
|
'--port', str(args_pytest["port"]),
|
|
'--extra-model-paths-config', 'tests/inference/extra_model_paths.yaml',
|
|
]
|
|
use_lru, lru_size = request.param
|
|
if use_lru:
|
|
pargs += ['--cache-lru', str(lru_size)]
|
|
print("Running server with args:", pargs) # noqa: T201
|
|
p = subprocess.Popen(pargs)
|
|
yield
|
|
p.kill()
|
|
torch.cuda.empty_cache()
|
|
|
|
def start_client(self, listen:str, port:int):
|
|
# Start client
|
|
comfy_client = ComfyClient()
|
|
# Connect to server (with retries)
|
|
n_tries = 5
|
|
for i in range(n_tries):
|
|
time.sleep(4)
|
|
try:
|
|
comfy_client.connect(listen=listen, port=port)
|
|
except ConnectionRefusedError as e:
|
|
print(e) # noqa: T201
|
|
print(f"({i+1}/{n_tries}) Retrying...") # noqa: T201
|
|
else:
|
|
break
|
|
return comfy_client
|
|
|
|
@fixture(scope="class", autouse=True)
|
|
def shared_client(self, args_pytest, _server):
|
|
client = self.start_client(args_pytest["listen"], args_pytest["port"])
|
|
yield client
|
|
del client
|
|
torch.cuda.empty_cache()
|
|
|
|
@fixture
|
|
def client(self, shared_client, request):
|
|
shared_client.set_test_name(f"execution[{request.node.name}]")
|
|
yield shared_client
|
|
|
|
@fixture
|
|
def builder(self, request):
|
|
yield GraphBuilder(prefix=request.node.name)
|
|
|
|
def test_lazy_input(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
|
|
mask = g.node("StubMask", value=0.0, height=512, width=512, batch_size=1)
|
|
|
|
lazy_mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))
|
|
output = g.node("SaveImage", images=lazy_mix.out(0))
|
|
result = client.run(g)
|
|
|
|
result_image = result.get_images(output)[0]
|
|
assert numpy.array(result_image).any() == 0, "Image should be black"
|
|
assert result.did_run(input1)
|
|
assert not result.did_run(input2)
|
|
assert result.did_run(mask)
|
|
assert result.did_run(lazy_mix)
|
|
|
|
def test_full_cache(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
input2 = g.node("StubImage", content="NOISE", height=512, width=512, batch_size=1)
|
|
mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)
|
|
|
|
lazy_mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))
|
|
g.node("SaveImage", images=lazy_mix.out(0))
|
|
|
|
client.run(g)
|
|
result2 = client.run(g)
|
|
for node_id, node in g.nodes.items():
|
|
assert not result2.did_run(node), f"Node {node_id} ran, but should have been cached"
|
|
|
|
def test_partial_cache(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
input2 = g.node("StubImage", content="NOISE", height=512, width=512, batch_size=1)
|
|
mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)
|
|
|
|
lazy_mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))
|
|
g.node("SaveImage", images=lazy_mix.out(0))
|
|
|
|
client.run(g)
|
|
mask.inputs['value'] = 0.4
|
|
result2 = client.run(g)
|
|
assert not result2.did_run(input1), "Input1 should have been cached"
|
|
assert not result2.did_run(input2), "Input2 should have been cached"
|
|
|
|
def test_error(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
# Different size of the two images
|
|
input2 = g.node("StubImage", content="NOISE", height=256, width=256, batch_size=1)
|
|
mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)
|
|
|
|
lazy_mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))
|
|
g.node("SaveImage", images=lazy_mix.out(0))
|
|
|
|
try:
|
|
client.run(g)
|
|
assert False, "Should have raised an error"
|
|
except Exception as e:
|
|
assert 'prompt_id' in e.args[0], f"Did not get back a proper error message: {e}"
|
|
|
|
@pytest.mark.parametrize("test_value, expect_error", [
|
|
(5, True),
|
|
("foo", True),
|
|
(5.0, False),
|
|
])
|
|
def test_validation_error_literal(self, test_value, expect_error, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
validation1 = g.node("TestCustomValidation1", input1=test_value, input2=3.0)
|
|
g.node("SaveImage", images=validation1.out(0))
|
|
|
|
if expect_error:
|
|
with pytest.raises(urllib.error.HTTPError):
|
|
client.run(g)
|
|
else:
|
|
client.run(g)
|
|
|
|
@pytest.mark.parametrize("test_type, test_value", [
|
|
("StubInt", 5),
|
|
("StubMask", 5.0)
|
|
])
|
|
def test_validation_error_edge1(self, test_type, test_value, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
stub = g.node(test_type, value=test_value)
|
|
validation1 = g.node("TestCustomValidation1", input1=stub.out(0), input2=3.0)
|
|
g.node("SaveImage", images=validation1.out(0))
|
|
|
|
with pytest.raises(urllib.error.HTTPError):
|
|
client.run(g)
|
|
|
|
@pytest.mark.parametrize("test_type, test_value, expect_error", [
|
|
("StubInt", 5, True),
|
|
("StubFloat", 5.0, False)
|
|
])
|
|
def test_validation_error_edge2(self, test_type, test_value, expect_error, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
stub = g.node(test_type, value=test_value)
|
|
validation2 = g.node("TestCustomValidation2", input1=stub.out(0), input2=3.0)
|
|
g.node("SaveImage", images=validation2.out(0))
|
|
|
|
if expect_error:
|
|
with pytest.raises(urllib.error.HTTPError):
|
|
client.run(g)
|
|
else:
|
|
client.run(g)
|
|
|
|
@pytest.mark.parametrize("test_type, test_value, expect_error", [
|
|
("StubInt", 5, True),
|
|
("StubFloat", 5.0, False)
|
|
])
|
|
def test_validation_error_edge3(self, test_type, test_value, expect_error, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
stub = g.node(test_type, value=test_value)
|
|
validation3 = g.node("TestCustomValidation3", input1=stub.out(0), input2=3.0)
|
|
g.node("SaveImage", images=validation3.out(0))
|
|
|
|
if expect_error:
|
|
with pytest.raises(urllib.error.HTTPError):
|
|
client.run(g)
|
|
else:
|
|
client.run(g)
|
|
|
|
@pytest.mark.parametrize("test_type, test_value, expect_error", [
|
|
("StubInt", 5, True),
|
|
("StubFloat", 5.0, False)
|
|
])
|
|
def test_validation_error_edge4(self, test_type, test_value, expect_error, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
stub = g.node(test_type, value=test_value)
|
|
validation4 = g.node("TestCustomValidation4", input1=stub.out(0), input2=3.0)
|
|
g.node("SaveImage", images=validation4.out(0))
|
|
|
|
if expect_error:
|
|
with pytest.raises(urllib.error.HTTPError):
|
|
client.run(g)
|
|
else:
|
|
client.run(g)
|
|
|
|
@pytest.mark.parametrize("test_value1, test_value2, expect_error", [
|
|
(0.0, 0.5, False),
|
|
(0.0, 5.0, False),
|
|
(0.0, 7.0, True)
|
|
])
|
|
def test_validation_error_kwargs(self, test_value1, test_value2, expect_error, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
validation5 = g.node("TestCustomValidation5", input1=test_value1, input2=test_value2)
|
|
g.node("SaveImage", images=validation5.out(0))
|
|
|
|
if expect_error:
|
|
with pytest.raises(urllib.error.HTTPError):
|
|
client.run(g)
|
|
else:
|
|
client.run(g)
|
|
|
|
def test_cycle_error(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
|
|
mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)
|
|
|
|
lazy_mix1 = g.node("TestLazyMixImages", image1=input1.out(0), mask=mask.out(0))
|
|
lazy_mix2 = g.node("TestLazyMixImages", image1=lazy_mix1.out(0), image2=input2.out(0), mask=mask.out(0))
|
|
g.node("SaveImage", images=lazy_mix2.out(0))
|
|
|
|
# When the cycle exists on initial submission, it should raise a validation error
|
|
with pytest.raises(urllib.error.HTTPError):
|
|
client.run(g)
|
|
|
|
def test_dynamic_cycle_error(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
|
|
generator = g.node("TestDynamicDependencyCycle", input1=input1.out(0), input2=input2.out(0))
|
|
g.node("SaveImage", images=generator.out(0))
|
|
|
|
# When the cycle is in a graph that is generated dynamically, it should raise a runtime error
|
|
try:
|
|
client.run(g)
|
|
assert False, "Should have raised an error"
|
|
except Exception as e:
|
|
assert 'prompt_id' in e.args[0], f"Did not get back a proper error message: {e}"
|
|
assert e.args[0]['node_id'] == generator.id, "Error should have been on the generator node"
|
|
|
|
def test_missing_node_error(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
input2 = g.node("StubImage", id="removeme", content="WHITE", height=512, width=512, batch_size=1)
|
|
input3 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
|
|
mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)
|
|
mix1 = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))
|
|
mix2 = g.node("TestLazyMixImages", image1=input1.out(0), image2=input3.out(0), mask=mask.out(0))
|
|
# We have multiple outputs. The first is invalid, but the second is valid
|
|
g.node("SaveImage", images=mix1.out(0))
|
|
g.node("SaveImage", images=mix2.out(0))
|
|
g.remove_node("removeme")
|
|
|
|
client.run(g)
|
|
|
|
# Add back in the missing node to make sure the error doesn't break the server
|
|
input2 = g.node("StubImage", id="removeme", content="WHITE", height=512, width=512, batch_size=1)
|
|
client.run(g)
|
|
|
|
def test_custom_is_changed(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
# Creating the nodes in this specific order previously caused a bug
|
|
save = g.node("SaveImage")
|
|
is_changed = g.node("TestCustomIsChanged", should_change=False)
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
|
|
save.set_input('images', is_changed.out(0))
|
|
is_changed.set_input('image', input1.out(0))
|
|
|
|
result1 = client.run(g)
|
|
result2 = client.run(g)
|
|
is_changed.set_input('should_change', True)
|
|
result3 = client.run(g)
|
|
result4 = client.run(g)
|
|
assert result1.did_run(is_changed), "is_changed should have been run"
|
|
assert not result2.did_run(is_changed), "is_changed should have been cached"
|
|
assert result3.did_run(is_changed), "is_changed should have been re-run"
|
|
assert result4.did_run(is_changed), "is_changed should not have been cached"
|
|
|
|
def test_undeclared_inputs(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
|
|
input3 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
input4 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
average = g.node("TestVariadicAverage", input1=input1.out(0), input2=input2.out(0), input3=input3.out(0), input4=input4.out(0))
|
|
output = g.node("SaveImage", images=average.out(0))
|
|
|
|
result = client.run(g)
|
|
result_image = result.get_images(output)[0]
|
|
expected = 255 // 4
|
|
assert numpy.array(result_image).min() == expected and numpy.array(result_image).max() == expected, "Image should be grey"
|
|
|
|
def test_for_loop(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
iterations = 4
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
|
|
is_changed = g.node("TestCustomIsChanged", should_change=True, image=input2.out(0))
|
|
for_open = g.node("TestForLoopOpen", remaining=iterations, initial_value1=is_changed.out(0))
|
|
average = g.node("TestVariadicAverage", input1=input1.out(0), input2=for_open.out(2))
|
|
for_close = g.node("TestForLoopClose", flow_control=for_open.out(0), initial_value1=average.out(0))
|
|
output = g.node("SaveImage", images=for_close.out(0))
|
|
|
|
for iterations in range(1, 5):
|
|
for_open.set_input('remaining', iterations)
|
|
result = client.run(g)
|
|
result_image = result.get_images(output)[0]
|
|
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"
|
|
|
|
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)
|
|
mask = g.node("StubMask", value=val_list.out(0), height=512, width=512, batch_size=1)
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
|
|
mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))
|
|
rebatch = g.node("RebatchImages", images=mix.out(0), batch_size=3)
|
|
output = g.node("SaveImage", images=rebatch.out(0))
|
|
|
|
result = client.run(g)
|
|
images = result.get_images(output)
|
|
assert len(images) == 3, "Should have 3 image"
|
|
assert numpy.array(images[0]).min() == 0 and numpy.array(images[0]).max() == 0, "First image should be 0.0"
|
|
assert numpy.array(images[1]).min() == 127 and numpy.array(images[1]).max() == 127, "Second image should be 0.5"
|
|
assert numpy.array(images[2]).min() == 255 and numpy.array(images[2]).max() == 255, "Third image should be 1.0"
|
|
|
|
def test_output_reuse(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
|
|
output1 = g.node("SaveImage", images=input1.out(0))
|
|
output2 = g.node("SaveImage", images=input1.out(0))
|
|
|
|
result = client.run(g)
|
|
images1 = result.get_images(output1)
|
|
images2 = result.get_images(output2)
|
|
assert len(images1) == 1, "Should have 1 image"
|
|
assert len(images2) == 1, "Should have 1 image"
|
|
|
|
|
|
# This tests that only constant outputs are used in the call to `IS_CHANGED`
|
|
def test_is_changed_with_outputs(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
input1 = g.node("StubConstantImage", value=0.5, height=512, width=512, batch_size=1)
|
|
test_node = g.node("TestIsChangedWithConstants", image=input1.out(0), value=0.5)
|
|
|
|
output = g.node("PreviewImage", images=test_node.out(0))
|
|
|
|
result = client.run(g)
|
|
images = result.get_images(output)
|
|
assert len(images) == 1, "Should have 1 image"
|
|
assert numpy.array(images[0]).min() == 63 and numpy.array(images[0]).max() == 63, "Image should have value 0.25"
|
|
|
|
result = client.run(g)
|
|
images = result.get_images(output)
|
|
assert len(images) == 1, "Should have 1 image"
|
|
assert numpy.array(images[0]).min() == 63 and numpy.array(images[0]).max() == 63, "Image should have value 0.25"
|
|
assert not result.did_run(test_node), "The execution should have been cached"
|
|
|
|
def test_parallel_sleep_nodes(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
image = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
|
|
# Create sleep nodes for each duration
|
|
sleep_node1 = g.node("TestSleep", value=image.out(0), seconds=2.8)
|
|
sleep_node2 = g.node("TestSleep", value=image.out(0), seconds=2.9)
|
|
sleep_node3 = g.node("TestSleep", value=image.out(0), seconds=3.0)
|
|
|
|
# Add outputs to verify the execution
|
|
_output1 = g.node("PreviewImage", images=sleep_node1.out(0))
|
|
_output2 = g.node("PreviewImage", images=sleep_node2.out(0))
|
|
_output3 = g.node("PreviewImage", images=sleep_node3.out(0))
|
|
|
|
start_time = time.time()
|
|
result = client.run(g)
|
|
elapsed_time = time.time() - start_time
|
|
|
|
# The test should take around 0.4 seconds (the longest sleep duration)
|
|
# plus some overhead, but definitely less than the sum of all sleeps (0.9s)
|
|
# We'll allow for up to 0.8s total to account for overhead
|
|
assert elapsed_time < 4.0, f"Parallel execution took {elapsed_time}s, expected less than 0.8s"
|
|
|
|
# Verify that all nodes executed
|
|
assert result.did_run(sleep_node1), "Sleep node 1 should have run"
|
|
assert result.did_run(sleep_node2), "Sleep node 2 should have run"
|
|
assert result.did_run(sleep_node3), "Sleep node 3 should have run"
|
|
|
|
def test_parallel_sleep_expansion(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
# Create input images with different values
|
|
image1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
image2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
|
|
image3 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
|
|
|
|
# Create a TestParallelSleep node that expands into multiple TestSleep nodes
|
|
parallel_sleep = g.node("TestParallelSleep",
|
|
image1=image1.out(0),
|
|
image2=image2.out(0),
|
|
image3=image3.out(0),
|
|
sleep1=0.4,
|
|
sleep2=0.5,
|
|
sleep3=0.6)
|
|
output = g.node("SaveImage", images=parallel_sleep.out(0))
|
|
|
|
start_time = time.time()
|
|
result = client.run(g)
|
|
elapsed_time = time.time() - start_time
|
|
|
|
# Similar to the previous test, expect parallel execution of the sleep nodes
|
|
# which should complete in less than the sum of all sleeps
|
|
assert elapsed_time < 0.8, f"Expansion execution took {elapsed_time}s, expected less than 0.8s"
|
|
|
|
# Verify the parallel sleep node executed
|
|
assert result.did_run(parallel_sleep), "ParallelSleep node should have run"
|
|
|
|
# Verify we get an image as output (blend of the three input images)
|
|
result_images = result.get_images(output)
|
|
assert len(result_images) == 1, "Should have 1 image"
|
|
# Average pixel value should be around 170 (255 * 2 // 3)
|
|
avg_value = numpy.array(result_images[0]).mean()
|
|
assert avg_value == 170, f"Image average value {avg_value} should be 170"
|
|
|
|
# This tests that nodes with OUTPUT_IS_LIST function correctly when they receive an ExecutionBlocker
|
|
# as input. We also test that when that list (containing an ExecutionBlocker) is passed to a node,
|
|
# only that one entry in the list is blocked.
|
|
def test_execution_block_list_output(self, client: ComfyClient, builder: GraphBuilder):
|
|
g = builder
|
|
image1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
image2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
|
|
image3 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
image_list = g.node("TestMakeListNode", value1=image1.out(0), value2=image2.out(0), value3=image3.out(0))
|
|
int1 = g.node("StubInt", value=1)
|
|
int2 = g.node("StubInt", value=2)
|
|
int3 = g.node("StubInt", value=3)
|
|
int_list = g.node("TestMakeListNode", value1=int1.out(0), value2=int2.out(0), value3=int3.out(0))
|
|
compare = g.node("TestIntConditions", a=int_list.out(0), b=2, operation="==")
|
|
blocker = g.node("TestExecutionBlocker", input=image_list.out(0), block=compare.out(0), verbose=False)
|
|
|
|
list_output = g.node("TestMakeListNode", value1=blocker.out(0))
|
|
output = g.node("PreviewImage", images=list_output.out(0))
|
|
|
|
result = client.run(g)
|
|
assert result.did_run(output), "The execution should have run"
|
|
images = result.get_images(output)
|
|
assert len(images) == 2, "Should have 2 images"
|
|
assert numpy.array(images[0]).min() == 0 and numpy.array(images[0]).max() == 0, "First image should be black"
|
|
assert numpy.array(images[1]).min() == 0 and numpy.array(images[1]).max() == 0, "Second image should also be black"
|
|
|
|
def test_ordered_history_endpoint(self, client: ComfyClient, builder: GraphBuilder):
|
|
"""Test the ordered history endpoint returns data in chronological order."""
|
|
# Clear history to start fresh
|
|
client.clear_history()
|
|
|
|
# Run multiple prompts to test ordering
|
|
prompt_ids = []
|
|
for _ in range(3):
|
|
g = builder
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
g.node("SaveImage", images=input1.out(0))
|
|
|
|
result = client.run(g)
|
|
prompt_ids.append(result.get_prompt_id())
|
|
time.sleep(0.1) # Small delay to ensure different timestamps
|
|
|
|
# Test ordered history endpoint
|
|
ordered_history = client.get_ordered_history()
|
|
assert "history" in ordered_history, "Ordered history should have history key"
|
|
assert isinstance(ordered_history["history"], list), "Ordered history should be a list"
|
|
assert len(ordered_history["history"]) == 3, "Should have exactly 3 prompts in history"
|
|
|
|
# Verify chronological ordering (most recent first)
|
|
history_prompt_ids = [item["prompt_id"] for item in ordered_history["history"]]
|
|
|
|
# Should be in chronological order (oldest first, as they're added in completion order)
|
|
assert history_prompt_ids == prompt_ids, "History should be in chronological order"
|
|
|
|
def test_history_prompt_id_endpoint(self, client: ComfyClient, builder: GraphBuilder):
|
|
"""Test fetching specific prompt history by ID."""
|
|
g = builder
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
g.node("SaveImage", images=input1.out(0))
|
|
|
|
result = client.run(g)
|
|
prompt_id = result.get_prompt_id()
|
|
|
|
# Test legacy history endpoint for specific prompt
|
|
legacy_history = client.get_history(prompt_id)
|
|
assert prompt_id in legacy_history, "Legacy history should contain prompt ID"
|
|
assert "outputs" in legacy_history[prompt_id], "Legacy history should have outputs"
|
|
|
|
# Test history_v2 endpoint for specific prompt
|
|
specific_history = client.get_history_v2_for_prompt(prompt_id)
|
|
assert prompt_id in specific_history, "History v2 should contain prompt ID"
|
|
assert specific_history[prompt_id] == legacy_history[prompt_id], "History v2 data should match legacy history"
|
|
|
|
def test_history_max_items(self, client: ComfyClient, builder: GraphBuilder):
|
|
"""Test legacy history endpoint with max_items parameter."""
|
|
# Clear history to start fresh
|
|
client.clear_history()
|
|
|
|
# Run multiple prompts to test pagination
|
|
for _ in range(5):
|
|
g = GraphBuilder() # Create fresh GraphBuilder for each run
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
g.node("SaveImage", images=input1.out(0))
|
|
|
|
client.run(g)
|
|
time.sleep(0.1) # Small delay to ensure different timestamps
|
|
|
|
# Test max_items parameter on legacy history endpoint
|
|
limited_history = client.get_history(max_items=2)
|
|
assert len(limited_history) == 2, "History should return exactly max_items"
|
|
|
|
def test_ordered_history_max_items_and_offset(self, client: ComfyClient, builder: GraphBuilder):
|
|
"""Test ordered history endpoint with max_items and offset parameters."""
|
|
# Clear history to start fresh
|
|
client.clear_history()
|
|
|
|
# Run multiple prompts to test pagination
|
|
for _ in range(5):
|
|
g = GraphBuilder() # Create fresh GraphBuilder for each run
|
|
input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
|
g.node("SaveImage", images=input1.out(0))
|
|
|
|
client.run(g)
|
|
time.sleep(0.1) # Small delay to ensure different timestamps
|
|
|
|
# Test max_items parameter on ordered history endpoint
|
|
limited_ordered = client.get_ordered_history(max_items=3)
|
|
assert "history" in limited_ordered, "Limited ordered history should have history key"
|
|
assert len(limited_ordered["history"]) == 3, "Ordered history should return exactly max_items"
|
|
|
|
# Test pagination with offset as cursor
|
|
full_history = client.get_ordered_history()
|
|
assert len(full_history["history"]) == 5, "Should have 5 items in full history"
|
|
|
|
# Extract prompt IDs from full history for comparison
|
|
full_prompt_ids = [item["prompt_id"] for item in full_history["history"]]
|
|
|
|
# Test proper pagination behavior with offset as cursor
|
|
|
|
# Test first page: offset=0, max_items=2
|
|
page1 = client.get_ordered_history(max_items=2, offset=0)
|
|
assert len(page1["history"]) == 2, "First page should have 2 items"
|
|
page1_ids = [item["prompt_id"] for item in page1["history"]]
|
|
assert page1_ids == full_prompt_ids[0:2], "First page should contain items at indices 0-1"
|
|
|
|
# Test second page: offset=2, max_items=2
|
|
page2 = client.get_ordered_history(max_items=2, offset=2)
|
|
assert len(page2["history"]) == 2, "Second page should have 2 items"
|
|
page2_ids = [item["prompt_id"] for item in page2["history"]]
|
|
assert page2_ids == full_prompt_ids[2:4], "Second page should contain items at indices 2-3"
|
|
|
|
# Test third page: offset=4, max_items=2
|
|
page3 = client.get_ordered_history(max_items=2, offset=4)
|
|
assert len(page3["history"]) == 1, "Third page should have 1 remaining item"
|
|
page3_ids = [item["prompt_id"] for item in page3["history"]]
|
|
assert page3_ids == full_prompt_ids[4:5], "Third page should contain item at index 4"
|
|
|
|
# Verify no overlap between pages
|
|
all_paginated_ids = page1_ids + page2_ids + page3_ids
|
|
assert len(set(all_paginated_ids)) == 5, "All paginated IDs should be unique"
|
|
assert set(all_paginated_ids) == set(full_prompt_ids), "Paginated results should cover all items"
|
|
|
|
# Test default behavior: get last N items (no offset specified)
|
|
# When offset < 0 and max_items is specified, offset = len(history) - max_items
|
|
last_2_items = client.get_ordered_history(max_items=2)
|
|
assert len(last_2_items["history"]) == 2, "Default behavior should return 2 items"
|
|
last_2_ids = [item["prompt_id"] for item in last_2_items["history"]]
|
|
# This should be equivalent to offset=3 (5-2=3)
|
|
assert last_2_ids == full_prompt_ids[3:5], "Default behavior should return last 2 items"
|
|
|
|
# Test offset beyond available items
|
|
beyond_offset = client.get_ordered_history(max_items=2, offset=10)
|
|
assert len(beyond_offset["history"]) == 0, "Offset beyond items should return empty list"
|