nef7/my-comfyui-workflow
0
1from io import BytesIO2import numpy3from PIL import Image4import pytest5from pytest import fixture6import time7import torch8from typing import Union, Dict9import json10import subprocess11import websocket #NOTE: websocket-client (https://github.com/websocket-client/websocket-client)12import uuid13import urllib.request14import urllib.parse15import urllib.error16from comfy_execution.graph_utils import GraphBuilder, Node17 18class RunResult:19 def __init__(self, prompt_id: str):20 self.outputs: Dict[str,Dict] = {}21 self.runs: Dict[str,bool] = {}22 self.prompt_id: str = prompt_id23 24 def get_output(self, node: Node):25 return self.outputs.get(node.id, None)26 27 def did_run(self, node: Node):28 return self.runs.get(node.id, False)29 30 def get_images(self, node: Node):31 output = self.get_output(node)32 if output is None:33 return []34 return output.get('image_objects', [])35 36 def get_prompt_id(self):37 return self.prompt_id38 39class ComfyClient:40 def __init__(self):41 self.test_name = ""42 43 def connect(self,44 listen:str = '127.0.0.1',45 port:Union[str,int] = 8188,46 client_id: str = str(uuid.uuid4())47 ):48 self.client_id = client_id49 self.server_address = f"{listen}:{port}"50 ws = websocket.WebSocket()51 ws.connect("ws://{}/ws?clientId={}".format(self.server_address, self.client_id))52 self.ws = ws53 54 def queue_prompt(self, prompt):55 p = {"prompt": prompt, "client_id": self.client_id}56 data = json.dumps(p).encode('utf-8')57 req = urllib.request.Request("http://{}/prompt".format(self.server_address), data=data)58 return json.loads(urllib.request.urlopen(req).read())59 60 def get_image(self, filename, subfolder, folder_type):61 data = {"filename": filename, "subfolder": subfolder, "type": folder_type}62 url_values = urllib.parse.urlencode(data)63 with urllib.request.urlopen("http://{}/view?{}".format(self.server_address, url_values)) as response:64 return response.read()65 66 def get_history(self, prompt_id):67 with urllib.request.urlopen("http://{}/history/{}".format(self.server_address, prompt_id)) as response:68 return json.loads(response.read())69 70 def set_test_name(self, name):71 self.test_name = name72 73 def run(self, graph):74 prompt = graph.finalize()75 for node in graph.nodes.values():76 if node.class_type == 'SaveImage':77 node.inputs['filename_prefix'] = self.test_name78 79 prompt_id = self.queue_prompt(prompt)['prompt_id']80 result = RunResult(prompt_id)81 while True:82 out = self.ws.recv()83 if isinstance(out, str):84 message = json.loads(out)85 if message['type'] == 'executing':86 data = message['data']87 if data['prompt_id'] != prompt_id:88 continue89 if data['node'] is None:90 break91 result.runs[data['node']] = True92 elif message['type'] == 'execution_error':93 raise Exception(message['data'])94 elif message['type'] == 'execution_cached':95 pass # Probably want to store this off for testing96 97 history = self.get_history(prompt_id)[prompt_id]98 for node_id in history['outputs']:99 node_output = history['outputs'][node_id]100 result.outputs[node_id] = node_output101 images_output = []102 if 'images' in node_output:103 for image in node_output['images']:104 image_data = self.get_image(image['filename'], image['subfolder'], image['type'])105 image_obj = Image.open(BytesIO(image_data))106 images_output.append(image_obj)107 node_output['image_objects'] = images_output108 109 return result110 111#112# Loop through these variables113#114@pytest.mark.execution115class TestExecution:116 #117 # Initialize server and client118 #119 @fixture(scope="class", autouse=True, params=[120 # (use_lru, lru_size)121 (False, 0),122 (True, 0),123 (True, 100),124 ])125 def _server(self, args_pytest, request):126 # Start server127 pargs = [128 'python','main.py',129 '--output-directory', args_pytest["output_dir"],130 '--listen', args_pytest["listen"],131 '--port', str(args_pytest["port"]),132 '--extra-model-paths-config', 'tests/inference/extra_model_paths.yaml',133 ]134 use_lru, lru_size = request.param135 if use_lru:136 pargs += ['--cache-lru', str(lru_size)]137 print("Running server with args:", pargs) # noqa: T201138 p = subprocess.Popen(pargs)139 yield140 p.kill()141 torch.cuda.empty_cache()142 143 def start_client(self, listen:str, port:int):144 # Start client145 comfy_client = ComfyClient()146 # Connect to server (with retries)147 n_tries = 5148 for i in range(n_tries):149 time.sleep(4)150 try:151 comfy_client.connect(listen=listen, port=port)152 except ConnectionRefusedError as e:153 print(e) # noqa: T201154 print(f"({i+1}/{n_tries}) Retrying...") # noqa: T201155 else:156 break157 return comfy_client158 159 @fixture(scope="class", autouse=True)160 def shared_client(self, args_pytest, _server):161 client = self.start_client(args_pytest["listen"], args_pytest["port"])162 yield client163 del client164 torch.cuda.empty_cache()165 166 @fixture167 def client(self, shared_client, request):168 shared_client.set_test_name(f"execution[{request.node.name}]")169 yield shared_client170 171 @fixture172 def builder(self, request):173 yield GraphBuilder(prefix=request.node.name)174 175 def test_lazy_input(self, client: ComfyClient, builder: GraphBuilder):176 g = builder177 input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)178 input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)179 mask = g.node("StubMask", value=0.0, height=512, width=512, batch_size=1)180 181 lazy_mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))182 output = g.node("SaveImage", images=lazy_mix.out(0))183 result = client.run(g)184 185 result_image = result.get_images(output)[0]186 assert numpy.array(result_image).any() == 0, "Image should be black"187 assert result.did_run(input1)188 assert not result.did_run(input2)189 assert result.did_run(mask)190 assert result.did_run(lazy_mix)191 192 def test_full_cache(self, client: ComfyClient, builder: GraphBuilder):193 g = builder194 input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)195 input2 = g.node("StubImage", content="NOISE", height=512, width=512, batch_size=1)196 mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)197 198 lazy_mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))199 g.node("SaveImage", images=lazy_mix.out(0))200 201 client.run(g)202 result2 = client.run(g)203 for node_id, node in g.nodes.items():204 assert not result2.did_run(node), f"Node {node_id} ran, but should have been cached"205 206 def test_partial_cache(self, client: ComfyClient, builder: GraphBuilder):207 g = builder208 input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)209 input2 = g.node("StubImage", content="NOISE", height=512, width=512, batch_size=1)210 mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)211 212 lazy_mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))213 g.node("SaveImage", images=lazy_mix.out(0))214 215 client.run(g)216 mask.inputs['value'] = 0.4217 result2 = client.run(g)218 assert not result2.did_run(input1), "Input1 should have been cached"219 assert not result2.did_run(input2), "Input2 should have been cached"220 221 def test_error(self, client: ComfyClient, builder: GraphBuilder):222 g = builder223 input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)224 # Different size of the two images225 input2 = g.node("StubImage", content="NOISE", height=256, width=256, batch_size=1)226 mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)227 228 lazy_mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))229 g.node("SaveImage", images=lazy_mix.out(0))230 231 try:232 client.run(g)233 assert False, "Should have raised an error"234 except Exception as e:235 assert 'prompt_id' in e.args[0], f"Did not get back a proper error message: {e}"236 237 @pytest.mark.parametrize("test_value, expect_error", [238 (5, True),239 ("foo", True),240 (5.0, False),241 ])242 def test_validation_error_literal(self, test_value, expect_error, client: ComfyClient, builder: GraphBuilder):243 g = builder244 validation1 = g.node("TestCustomValidation1", input1=test_value, input2=3.0)245 g.node("SaveImage", images=validation1.out(0))246 247 if expect_error:248 with pytest.raises(urllib.error.HTTPError):249 client.run(g)250 else:251 client.run(g)252 253 @pytest.mark.parametrize("test_type, test_value", [254 ("StubInt", 5),255 ("StubFloat", 5.0)256 ])257 def test_validation_error_edge1(self, test_type, test_value, client: ComfyClient, builder: GraphBuilder):258 g = builder259 stub = g.node(test_type, value=test_value)260 validation1 = g.node("TestCustomValidation1", input1=stub.out(0), input2=3.0)261 g.node("SaveImage", images=validation1.out(0))262 263 with pytest.raises(urllib.error.HTTPError):264 client.run(g)265 266 @pytest.mark.parametrize("test_type, test_value, expect_error", [267 ("StubInt", 5, True),268 ("StubFloat", 5.0, False)269 ])270 def test_validation_error_edge2(self, test_type, test_value, expect_error, client: ComfyClient, builder: GraphBuilder):271 g = builder272 stub = g.node(test_type, value=test_value)273 validation2 = g.node("TestCustomValidation2", input1=stub.out(0), input2=3.0)274 g.node("SaveImage", images=validation2.out(0))275 276 if expect_error:277 with pytest.raises(urllib.error.HTTPError):278 client.run(g)279 else:280 client.run(g)281 282 @pytest.mark.parametrize("test_type, test_value, expect_error", [283 ("StubInt", 5, True),284 ("StubFloat", 5.0, False)285 ])286 def test_validation_error_edge3(self, test_type, test_value, expect_error, client: ComfyClient, builder: GraphBuilder):287 g = builder288 stub = g.node(test_type, value=test_value)289 validation3 = g.node("TestCustomValidation3", input1=stub.out(0), input2=3.0)290 g.node("SaveImage", images=validation3.out(0))291 292 if expect_error:293 with pytest.raises(urllib.error.HTTPError):294 client.run(g)295 else:296 client.run(g)297 298 @pytest.mark.parametrize("test_type, test_value, expect_error", [299 ("StubInt", 5, True),300 ("StubFloat", 5.0, False)301 ])302 def test_validation_error_edge4(self, test_type, test_value, expect_error, client: ComfyClient, builder: GraphBuilder):303 g = builder304 stub = g.node(test_type, value=test_value)305 validation4 = g.node("TestCustomValidation4", input1=stub.out(0), input2=3.0)306 g.node("SaveImage", images=validation4.out(0))307 308 if expect_error:309 with pytest.raises(urllib.error.HTTPError):310 client.run(g)311 else:312 client.run(g)313 314 @pytest.mark.parametrize("test_value1, test_value2, expect_error", [315 (0.0, 0.5, False),316 (0.0, 5.0, False),317 (0.0, 7.0, True)318 ])319 def test_validation_error_kwargs(self, test_value1, test_value2, expect_error, client: ComfyClient, builder: GraphBuilder):320 g = builder321 validation5 = g.node("TestCustomValidation5", input1=test_value1, input2=test_value2)322 g.node("SaveImage", images=validation5.out(0))323 324 if expect_error:325 with pytest.raises(urllib.error.HTTPError):326 client.run(g)327 else:328 client.run(g)329 330 def test_cycle_error(self, client: ComfyClient, builder: GraphBuilder):331 g = builder332 input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)333 input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)334 mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)335 336 lazy_mix1 = g.node("TestLazyMixImages", image1=input1.out(0), mask=mask.out(0))337 lazy_mix2 = g.node("TestLazyMixImages", image1=lazy_mix1.out(0), image2=input2.out(0), mask=mask.out(0))338 g.node("SaveImage", images=lazy_mix2.out(0))339 340 # When the cycle exists on initial submission, it should raise a validation error341 with pytest.raises(urllib.error.HTTPError):342 client.run(g)343 344 def test_dynamic_cycle_error(self, client: ComfyClient, builder: GraphBuilder):345 g = builder346 input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)347 input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)348 generator = g.node("TestDynamicDependencyCycle", input1=input1.out(0), input2=input2.out(0))349 g.node("SaveImage", images=generator.out(0))350 351 # When the cycle is in a graph that is generated dynamically, it should raise a runtime error352 try:353 client.run(g)354 assert False, "Should have raised an error"355 except Exception as e:356 assert 'prompt_id' in e.args[0], f"Did not get back a proper error message: {e}"357 assert e.args[0]['node_id'] == generator.id, "Error should have been on the generator node"358 359 def test_missing_node_error(self, client: ComfyClient, builder: GraphBuilder):360 g = builder361 input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)362 input2 = g.node("StubImage", id="removeme", content="WHITE", height=512, width=512, batch_size=1)363 input3 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)364 mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)365 mix1 = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))366 mix2 = g.node("TestLazyMixImages", image1=input1.out(0), image2=input3.out(0), mask=mask.out(0))367 # We have multiple outputs. The first is invalid, but the second is valid368 g.node("SaveImage", images=mix1.out(0))369 g.node("SaveImage", images=mix2.out(0))370 g.remove_node("removeme")371 372 client.run(g)373 374 # Add back in the missing node to make sure the error doesn't break the server375 input2 = g.node("StubImage", id="removeme", content="WHITE", height=512, width=512, batch_size=1)376 client.run(g)377 378 def test_custom_is_changed(self, client: ComfyClient, builder: GraphBuilder):379 g = builder380 # Creating the nodes in this specific order previously caused a bug381 save = g.node("SaveImage")382 is_changed = g.node("TestCustomIsChanged", should_change=False)383 input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)384 385 save.set_input('images', is_changed.out(0))386 is_changed.set_input('image', input1.out(0))387 388 result1 = client.run(g)389 result2 = client.run(g)390 is_changed.set_input('should_change', True)391 result3 = client.run(g)392 result4 = client.run(g)393 assert result1.did_run(is_changed), "is_changed should have been run"394 assert not result2.did_run(is_changed), "is_changed should have been cached"395 assert result3.did_run(is_changed), "is_changed should have been re-run"396 assert result4.did_run(is_changed), "is_changed should not have been cached"397 398 def test_undeclared_inputs(self, client: ComfyClient, builder: GraphBuilder):399 g = builder400 input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)401 input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)402 input3 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)403 input4 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)404 average = g.node("TestVariadicAverage", input1=input1.out(0), input2=input2.out(0), input3=input3.out(0), input4=input4.out(0))405 output = g.node("SaveImage", images=average.out(0))406 407 result = client.run(g)408 result_image = result.get_images(output)[0]409 expected = 255 // 4410 assert numpy.array(result_image).min() == expected and numpy.array(result_image).max() == expected, "Image should be grey"411 412 def test_for_loop(self, client: ComfyClient, builder: GraphBuilder):413 g = builder414 iterations = 4415 input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)416 input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)417 is_changed = g.node("TestCustomIsChanged", should_change=True, image=input2.out(0))418 for_open = g.node("TestForLoopOpen", remaining=iterations, initial_value1=is_changed.out(0))419 average = g.node("TestVariadicAverage", input1=input1.out(0), input2=for_open.out(2))420 for_close = g.node("TestForLoopClose", flow_control=for_open.out(0), initial_value1=average.out(0))421 output = g.node("SaveImage", images=for_close.out(0))422 423 for iterations in range(1, 5):424 for_open.set_input('remaining', iterations)425 result = client.run(g)426 result_image = result.get_images(output)[0]427 expected = 255 // (2 ** iterations)428 assert numpy.array(result_image).min() == expected and numpy.array(result_image).max() == expected, "Image should be grey"429 assert result.did_run(is_changed)430 431 def test_mixed_expansion_returns(self, client: ComfyClient, builder: GraphBuilder):432 g = builder433 val_list = g.node("TestMakeListNode", value1=0.1, value2=0.2, value3=0.3)434 mixed = g.node("TestMixedExpansionReturns", input1=val_list.out(0))435 output_dynamic = g.node("SaveImage", images=mixed.out(0))436 output_literal = g.node("SaveImage", images=mixed.out(1))437 438 result = client.run(g)439 images_dynamic = result.get_images(output_dynamic)440 assert len(images_dynamic) == 3, "Should have 2 images"441 assert numpy.array(images_dynamic[0]).min() == 25 and numpy.array(images_dynamic[0]).max() == 25, "First image should be 0.1"442 assert numpy.array(images_dynamic[1]).min() == 51 and numpy.array(images_dynamic[1]).max() == 51, "Second image should be 0.2"443 assert numpy.array(images_dynamic[2]).min() == 76 and numpy.array(images_dynamic[2]).max() == 76, "Third image should be 0.3"444 445 images_literal = result.get_images(output_literal)446 assert len(images_literal) == 3, "Should have 2 images"447 for i in range(3):448 assert numpy.array(images_literal[i]).min() == 255 and numpy.array(images_literal[i]).max() == 255, "All images should be white"449 450 def test_mixed_lazy_results(self, client: ComfyClient, builder: GraphBuilder):451 g = builder452 val_list = g.node("TestMakeListNode", value1=0.0, value2=0.5, value3=1.0)453 mask = g.node("StubMask", value=val_list.out(0), height=512, width=512, batch_size=1)454 input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)455 input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)456 mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))457 rebatch = g.node("RebatchImages", images=mix.out(0), batch_size=3)458 output = g.node("SaveImage", images=rebatch.out(0))459 460 result = client.run(g)461 images = result.get_images(output)462 assert len(images) == 3, "Should have 3 image"463 assert numpy.array(images[0]).min() == 0 and numpy.array(images[0]).max() == 0, "First image should be 0.0"464 assert numpy.array(images[1]).min() == 127 and numpy.array(images[1]).max() == 127, "Second image should be 0.5"465 assert numpy.array(images[2]).min() == 255 and numpy.array(images[2]).max() == 255, "Third image should be 1.0"466 467 def test_output_reuse(self, client: ComfyClient, builder: GraphBuilder):468 g = builder469 input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)470 471 output1 = g.node("SaveImage", images=input1.out(0))472 output2 = g.node("SaveImage", images=input1.out(0))473 474 result = client.run(g)475 images1 = result.get_images(output1)476 images2 = result.get_images(output2)477 assert len(images1) == 1, "Should have 1 image"478 assert len(images2) == 1, "Should have 1 image"479 480 481 # This tests that only constant outputs are used in the call to `IS_CHANGED`482 def test_is_changed_with_outputs(self, client: ComfyClient, builder: GraphBuilder):483 g = builder484 input1 = g.node("StubConstantImage", value=0.5, height=512, width=512, batch_size=1)485 test_node = g.node("TestIsChangedWithConstants", image=input1.out(0), value=0.5)486 487 output = g.node("PreviewImage", images=test_node.out(0))488 489 result = client.run(g)490 images = result.get_images(output)491 assert len(images) == 1, "Should have 1 image"492 assert numpy.array(images[0]).min() == 63 and numpy.array(images[0]).max() == 63, "Image should have value 0.25"493 494 result = client.run(g)495 images = result.get_images(output)496 assert len(images) == 1, "Should have 1 image"497 assert numpy.array(images[0]).min() == 63 and numpy.array(images[0]).max() == 63, "Image should have value 0.25"498 assert not result.did_run(test_node), "The execution should have been cached"499 500 # This tests that nodes with OUTPUT_IS_LIST function correctly when they receive an ExecutionBlocker501 # as input. We also test that when that list (containing an ExecutionBlocker) is passed to a node,502 # only that one entry in the list is blocked.503 def test_execution_block_list_output(self, client: ComfyClient, builder: GraphBuilder):504 g = builder505 image1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)506 image2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)507 image3 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)508 image_list = g.node("TestMakeListNode", value1=image1.out(0), value2=image2.out(0), value3=image3.out(0))509 int1 = g.node("StubInt", value=1)510 int2 = g.node("StubInt", value=2)511 int3 = g.node("StubInt", value=3)512 int_list = g.node("TestMakeListNode", value1=int1.out(0), value2=int2.out(0), value3=int3.out(0))513 compare = g.node("TestIntConditions", a=int_list.out(0), b=2, operation="==")514 blocker = g.node("TestExecutionBlocker", input=image_list.out(0), block=compare.out(0), verbose=False)515 516 list_output = g.node("TestMakeListNode", value1=blocker.out(0))517 output = g.node("PreviewImage", images=list_output.out(0))518 519 result = client.run(g)520 assert result.did_run(output), "The execution should have run"521 images = result.get_images(output)522 assert len(images) == 2, "Should have 2 images"523 assert numpy.array(images[0]).min() == 0 and numpy.array(images[0]).max() == 0, "First image should be black"524 assert numpy.array(images[1]).min() == 0 and numpy.array(images[1]).max() == 0, "Second image should also be black"525