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:
gramps
2026-07-27 08:06:03 -07:00
parent 576d9333b3
commit aecd3330fd
15 changed files with 2115 additions and 1005 deletions
+69
View File
@@ -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),
]