Improved API support

- Run comfyui workflows directly inside other python applications using
   EmbeddedComfyClient.
 - Optional telemetry in prompts and models using anonymity preserving
   Plausible self-hosted or hosted.
 - Better OpenAPI schema
 - Basic support for distributed ComfyUI backends. Limitations: no
   progress reporting, no easy way to start your own distributed
   backend, requires RabbitMQ as a message broker.

tests/distributed/test_asyncio_remote_client.py

@@ -0,0 +1,19 @@
+import pytest
+from comfy.client.aio_client import AsyncRemoteComfyClient
+from comfy.client.sdxl_with_refiner_workflow import sdxl_workflow_with_refiner
+
+
+@pytest.mark.asyncio
+async def test_completes_prompt(comfy_background_server):
+    client = AsyncRemoteComfyClient()
+    prompt = sdxl_workflow_with_refiner("test", inference_steps=1, refiner_steps=1)
+    png_image_bytes = await client.queue_prompt(prompt)
+    assert len(png_image_bytes) > 1000
+
+@pytest.mark.asyncio
+async def test_completes_prompt_with_ui(comfy_background_server):
+    client = AsyncRemoteComfyClient()
+    prompt = sdxl_workflow_with_refiner("test", inference_steps=1, refiner_steps=1)
+    result_dict = await client.queue_prompt_ui(prompt)
+    # should contain one output
+    assert len(result_dict) == 1

tests/distributed/test_distributed_queue.py

@@ -0,0 +1,39 @@
+import os
+import uuid
+
+import pytest
+
+from comfy.client.embedded_comfy_client import EmbeddedComfyClient, ServerStub
+from comfy.client.sdxl_with_refiner_workflow import sdxl_workflow_with_refiner
+from comfy.component_model.make_mutable import make_mutable
+from comfy.component_model.queue_types import QueueItem, QueueTuple, TaskInvocation
+from comfy.distributed.distributed_prompt_worker import DistributedPromptWorker
+
+
+@pytest.mark.asyncio
+async def test_basic_queue_worker() -> None:
+    os.environ["TC_HOST"] = "localhost"
+    # there are lots of side effects from importing that we have to deal with
+    from testcontainers.rabbitmq import RabbitMqContainer
+
+    with RabbitMqContainer("rabbitmq:latest") as rabbitmq:
+        params = rabbitmq.get_connection_params()
+        async with EmbeddedComfyClient() as client:
+            async with DistributedPromptWorker(client,
+                                               connection_uri=f"amqp://guest:guest@127.0.0.1:{params.port}") as worker:
+                # this unfortunately does a bunch of initialization on the test thread
+                from comfy.cmd.execution import validate_prompt
+                from comfy.distributed.distributed_prompt_queue import DistributedPromptQueue
+                # now submit some jobs
+                distributed_queue = DistributedPromptQueue(ServerStub(), is_callee=False, is_caller=True,
+                                                           connection_uri=f"amqp://guest:guest@127.0.0.1:{params.port}")
+                await distributed_queue.init()
+                prompt = make_mutable(sdxl_workflow_with_refiner("test", inference_steps=1, refiner_steps=1))
+                validation_tuple = validate_prompt(prompt)
+                item_id = str(uuid.uuid4())
+                queue_tuple: QueueTuple = (0, item_id, prompt, {}, validation_tuple[2])
+                res: TaskInvocation = await distributed_queue.put_async(QueueItem(queue_tuple, None))
+                assert res.item_id == item_id
+                assert len(res.outputs) == 1
+                assert res.status is not None
+                assert res.status.status_str == "success"

tests/distributed/test_embedded_client.py

@@ -0,0 +1,33 @@
+import pytest
+import torch
+
+from comfy.client.embedded_comfy_client import EmbeddedComfyClient
+from comfy.client.sdxl_with_refiner_workflow import sdxl_workflow_with_refiner
+
+
+@pytest.mark.asyncio
+async def test_cuda_memory_usage():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is not available in this environment")
+
+    device = torch.device("cuda")
+    starting_memory = torch.cuda.memory_allocated(device)
+
+    async with EmbeddedComfyClient() as client:
+        prompt = sdxl_workflow_with_refiner("test")
+        outputs = await client.queue_prompt(prompt)
+        assert outputs["13"]["images"][0]["abs_path"] is not None
+        memory_after_workflow = torch.cuda.memory_allocated(device)
+        assert memory_after_workflow > starting_memory, "Expected CUDA memory to increase after running the workflow"
+
+    ending_memory = torch.cuda.memory_allocated(device)
+    assert abs(
+        ending_memory - starting_memory) < 1e7, "Expected CUDA memory to return close to starting memory after cleanup"
+
+
+@pytest.mark.asyncio
+async def test_embedded_comfy():
+    async with EmbeddedComfyClient() as client:
+        prompt = sdxl_workflow_with_refiner("test")
+        outputs = await client.queue_prompt(prompt)
+        assert outputs["13"]["images"][0]["abs_path"] is not None