Files
cAIc/routers/image.py
T
gramps aecd3330fd 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
2026-07-27 08:06:03 -07:00

71 lines
2.2 KiB
Python

"""JarvisChat routers — Image generation proxy endpoint."""
import base64
import logging
from fastapi import APIRouter, HTTPException
from fastapi.responses import Response
from cluster import CLUSTER_NODES, request_image_generate
log = logging.getLogger("caic")
router = APIRouter()
def _find_image_node() -> str | None:
for name, node in CLUSTER_NODES.items():
if node.get("status") == "active" and "image_gen" in node.get("capabilities", []):
return name
return None
@router.post("/api/image/generate")
async def generate_image(request_body: dict):
prompt = (request_body.get("prompt") or "").strip()
if not prompt:
raise HTTPException(status_code=400, detail="Prompt is required")
negative_prompt = request_body.get("negative_prompt", "")
width = min(max(request_body.get("width", 1024), 256), 2048)
height = min(max(request_body.get("height", 1024), 256), 2048)
steps = min(max(request_body.get("steps", 20), 1), 50)
seed = request_body.get("seed", -1)
model = request_body.get("model", "")
node_name = _find_image_node()
if not node_name:
raise HTTPException(status_code=503, detail="No image generation service available")
log.info("image generate via %s: %s", node_name, prompt[:60])
image_b64 = await request_image_generate(
node_name=node_name,
prompt=prompt,
negative_prompt=negative_prompt,
width=width,
height=height,
steps=steps,
seed=seed,
model=model,
)
if image_b64 is None:
raise HTTPException(status_code=504, detail="Image generation timed out or failed")
image_bytes = base64.b64decode(image_b64)
return Response(content=image_bytes, media_type="image/png")
@router.get("/api/image/status")
async def image_status():
nodes = []
for name, node in CLUSTER_NODES.items():
caps = node.get("capabilities", [])
if "image_gen" in caps:
nodes.append({
"name": name,
"status": node.get("status"),
"load": node.get("load"),
"last_seen": node.get("last_seen"),
})
return {"available": len(nodes) > 0, "nodes": nodes}