mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-11 14:50:49 +08:00
1b2ea61345
- 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.
40 lines
2.1 KiB
Python
40 lines
2.1 KiB
Python
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"
|