feat: image generation service — ComfyUI cluster integration
- POST /api/image/generate proxy endpoint (admin required) - GET /api/image/status lists available image gen nodes - Cluster AMQP protocol: cmd.image_generate, image_generated, image_failed - Node agent auto-detects ComfyUI, registers image_gen capability - ComfyUI workflow builder: CheckpointLoader → KSampler → VAEDecode → SaveImage - Hardware probe checks ComfyUI reachability + checkpoint model list - 27 tests covering cluster handlers, router, node agent, hardware, capabilities - Config: CAIC_COMFYUI_BASE, CAIC_COMFYUI_TIMEOUT, comfyui_port in agent.ini - Version bump to v1.1.0 - Documentation: ai.md, wiki/Developer-Architecture.md, current-wip.md, README.md, .env.example
This commit is contained in:
+69
@@ -18,6 +18,7 @@ CLUSTER_NODES: dict[str, dict] = {}
|
||||
CLUSTER_EVENTS: deque = deque(maxlen=1000)
|
||||
CLUSTER_COORDINATOR: str | None = None
|
||||
_pending_pings: dict[str, tuple[str, asyncio.Event]] = {}
|
||||
_pending_image: dict[str, tuple[str, asyncio.Event]] = {}
|
||||
NODE_NAME: str = os.environ.get("CAIC_NODE_NAME", "ultron")
|
||||
PING_TIMEOUT: float = 5.0
|
||||
|
||||
@@ -240,6 +241,72 @@ async def handle_model_failed(exchange: str, routing_key: str, payload: dict) ->
|
||||
_push_event("cluster", "error", node_name, f"Model swap failed: {error}")
|
||||
|
||||
|
||||
async def handle_image_generated(exchange: str, routing_key: str, payload: dict) -> None:
|
||||
node_name = payload.get("node_name", routing_key.split(".")[1] if "." in routing_key else "unknown")
|
||||
request_id = payload.get("request_id")
|
||||
|
||||
if request_id and request_id in _pending_image:
|
||||
_, event = _pending_image.pop(request_id)
|
||||
_pending_image[request_id] = (payload.get("image_base64", ""), event)
|
||||
event.set()
|
||||
|
||||
if node_name in CLUSTER_NODES:
|
||||
CLUSTER_NODES[node_name]["last_seen"] = datetime.now(timezone.utc).isoformat() + "Z"
|
||||
|
||||
|
||||
async def handle_image_failed(exchange: str, routing_key: str, payload: dict) -> None:
|
||||
node_name = payload.get("node_name", routing_key.split(".")[1] if "." in routing_key else "unknown")
|
||||
request_id = payload.get("request_id")
|
||||
error = payload.get("error", "unknown error")
|
||||
|
||||
if request_id and request_id in _pending_image:
|
||||
_pending_image[request_id] = ("", _pending_image[request_id][1])
|
||||
_pending_image[request_id][1].set()
|
||||
|
||||
_push_event("application", "error", node_name, f"Image generation failed: {error}")
|
||||
|
||||
|
||||
async def request_image_generate(
|
||||
node_name: str, prompt: str, negative_prompt: str = "",
|
||||
width: int = 1024, height: int = 1024, steps: int = 20,
|
||||
seed: int = -1, model: str = "", timeout: float = 120,
|
||||
) -> str | None:
|
||||
if node_name not in CLUSTER_NODES:
|
||||
log.warning("request_image_generate: unknown node %s", node_name)
|
||||
return None
|
||||
|
||||
caps = CLUSTER_NODES[node_name].get("capabilities", [])
|
||||
if "image_gen" not in caps:
|
||||
log.warning("request_image_generate: node %s lacks image_gen capability", node_name)
|
||||
return None
|
||||
|
||||
request_id = str(uuid.uuid4())
|
||||
event = asyncio.Event()
|
||||
_pending_image[request_id] = ("", event)
|
||||
|
||||
now = datetime.now(timezone.utc).isoformat() + "Z"
|
||||
_push_event("application", "info", node_name, f"Image generation requested: {prompt[:60]}...")
|
||||
|
||||
await publish(AMQP_EXCHANGE_ADMIN, f"node.{node_name}.cmd.image_generate", {
|
||||
"from": NODE_NAME, "type": "image_generate",
|
||||
"request_id": request_id,
|
||||
"prompt": prompt,
|
||||
"negative_prompt": negative_prompt,
|
||||
"width": width, "height": height,
|
||||
"steps": steps, "seed": seed, "model": model,
|
||||
"timestamp": now,
|
||||
})
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(event.wait(), timeout=timeout)
|
||||
result = _pending_image.pop(request_id, (None, None))
|
||||
return result[0]
|
||||
except asyncio.TimeoutError:
|
||||
_pending_image.pop(request_id, None)
|
||||
_push_event("application", "warn", node_name, "Image generation timed out")
|
||||
return None
|
||||
|
||||
|
||||
SUBSCRIBE_TABLE = [
|
||||
(AMQP_EXCHANGE_ADMIN, ["node.*.register"], handle_registration),
|
||||
(AMQP_EXCHANGE_ADMIN, ["node.*.deregister"], handle_deregistration),
|
||||
@@ -249,6 +316,8 @@ SUBSCRIBE_TABLE = [
|
||||
(AMQP_EXCHANGE_SYSTEM, ["node.*.heartbeat"], handle_heartbeat),
|
||||
(AMQP_EXCHANGE_SYSTEM, ["node.*.model_ready"], handle_model_ready),
|
||||
(AMQP_EXCHANGE_SYSTEM, ["node.*.model_failed"], handle_model_failed),
|
||||
(AMQP_EXCHANGE_SYSTEM, ["node.*.image_generated"], handle_image_generated),
|
||||
(AMQP_EXCHANGE_SYSTEM, ["node.*.image_failed"], handle_image_failed),
|
||||
]
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user