gramps/cAIc
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: image generation service — ComfyUI cluster integration

Browse Source 
- 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
Download Patch File Download Diff File Expand all files Collapse all files
node_agent/agent.py
+171
View File
@@ -93,6 +93,7 @@ class AgentConfig:
self.capabilities: list[str] = ["llm"]
self.amqp_url: str = "amqp://caic:password@localhost:5672/caic"
self.llama_port: int = 8081
self.comfyui_port: int = 8188
self.models_dir: str = "/var/lib/caic/models"
self.active_model: str = ""
@@ -113,6 +114,7 @@ class AgentConfig:
cfg.capabilities = [c.strip() for c in raw_caps.split(",") if c.strip()]
cfg.amqp_url = parser.get(sec, "amqp_url", fallback=cfg.amqp_url)
cfg.llama_port = parser.getint(sec, "llama_port", fallback=cfg.llama_port)
cfg.comfyui_port = parser.getint(sec, "comfyui_port", fallback=cfg.comfyui_port)
cfg.models_dir = parser.get(sec, "models_dir", fallback=cfg.models_dir)
cfg.active_model = parser.get(sec, "active_model", fallback=cfg.active_model)
return cfg
@@ -212,6 +214,19 @@ def get_load() -> dict:
return load
def detect_capabilities(cfg: AgentConfig) -> list[str]:
caps = list(cfg.capabilities)
if "image_gen" not in caps and HAS_HTTPX:
try:
resp = httpx.get(f"http://localhost:{cfg.comfyui_port}/system_stats", timeout=3)
if resp.status_code == 200:
caps.append("image_gen")
log.info("auto-detected image_gen capability (ComfyUI on port %d)", cfg.comfyui_port)
except Exception:
pass
return caps
# ── AMQP helpers ────────────────────────────────────────────────────────
async def declare_exchanges(channel) -> tuple:
@@ -356,6 +371,152 @@ async def _wait_for_llama(port: int, timeout: int = 120, interval: int = 2) -> b
return False
# ── image generation ─────────────────────────────────────────────────────
async def handle_image_generate(cfg: AgentConfig, channel, exchanges, msg: aio_pika.IncomingMessage):
admin_ex, system_ex = exchanges
async with msg.process():
try:
payload = json.loads(msg.body.decode())
except json.JSONDecodeError:
return
request_id = payload.get("request_id")
prompt = payload.get("prompt", "")
negative_prompt = payload.get("negative_prompt", "")
width = payload.get("width", 1024)
height = payload.get("height", 1024)
steps = payload.get("steps", 20)
seed = payload.get("seed", -1)
model = payload.get("model", "")
if not prompt:
log.error("image_generate missing prompt")
return
log.info("image generate: prompt=%s %dx%d steps=%d", prompt[:60], width, height, steps)
now = datetime.now(timezone.utc).isoformat() + "Z"
try:
image_data = await _comfyui_generate(
cfg, prompt, negative_prompt, width, height, steps, seed, model,
)
result_payload = {
"node_name": cfg.node_name,
"type": "image_generated",
"request_id": request_id,
"image_base64": image_data,
"timestamp": now,
}
log.info("image generate complete: request_id=%s", request_id)
except Exception as e:
result_payload = {
"node_name": cfg.node_name,
"type": "image_failed",
"request_id": request_id,
"error": str(e),
"timestamp": now,
}
log.error("image generate failed: %s", e)
await publish(channel, system_ex, f"node.{cfg.node_name}.{result_payload['type']}", result_payload)
async def _comfyui_generate(
cfg: AgentConfig, prompt: str, negative_prompt: str,
width: int, height: int, steps: int, seed: int, model: str,
) -> str:
import random
import uuid as _uuid
if not HAS_HTTPX:
raise RuntimeError("httpx not installed")
client_id = str(_uuid.uuid4())
if seed < 0:
seed = random.randint(0, 2**32 - 1)
checkpoint = model or "model.safetensors"
workflow = {
"3": {
"class_type": "KSampler",
"inputs": {
"seed": seed,
"steps": steps,
"cfg": 7.0,
"sampler_name": "euler",
"scheduler": "normal",
"denoise": 1.0,
"model": ["4", 0],
"positive": ["6", 0],
"negative": ["7", 0],
"latent_image": ["5", 0],
},
},
"4": {
"class_type": "CheckpointLoaderSimple",
"inputs": {"ckpt_name": checkpoint},
},
"5": {
"class_type": "EmptyLatentImage",
"inputs": {"width": width, "height": height, "batch_size": 1},
},
"6": {
"class_type": "CLIPTextEncode",
"inputs": {"text": prompt, "clip": ["4", 1]},
},
"7": {
"class_type": "CLIPTextEncode",
"inputs": {"text": negative_prompt or "blurry, low quality", "clip": ["4", 1]},
},
"8": {
"class_type": "VAEDecode",
"inputs": {"samples": ["3", 0], "vae": ["4", 2]},
},
"9": {
"class_type": "SaveImage",
"inputs": {"filename_prefix": f"caic_{client_id}", "images": ["8", 0]},
},
}
async with httpx.AsyncClient(timeout=120) as client:
resp = await client.post(
f"http://localhost:{cfg.comfyui_port}/prompt",
json={"prompt": workflow, "client_id": client_id},
)
if resp.status_code != 200:
raise RuntimeError(f"ComfyUI prompt failed: {resp.status_code} {resp.text}")
prompt_id = resp.json().get("prompt_id")
if not prompt_id:
raise RuntimeError("ComfyUI returned no prompt_id")
deadline = time.time() + 120
while time.time() < deadline:
resp = await client.get(f"http://localhost:{cfg.comfyui_port}/history/{prompt_id}")
if resp.status_code == 200:
history = resp.json().get(prompt_id, {})
outputs = history.get("outputs", {})
for node_id, node_output in outputs.items():
images = node_output.get("images", [])
if images:
img_info = images[0]
filename = img_info.get("filename")
subfolder = img_info.get("subfolder", "")
img_type = img_info.get("type", "output")
img_resp = await client.get(
f"http://localhost:{cfg.comfyui_port}/view",
params={"filename": filename, "subfolder": subfolder, "type": img_type},
)
if img_resp.status_code == 200:
import base64
return base64.b64encode(img_resp.content).decode()
await asyncio.sleep(1)
raise RuntimeError("ComfyUI generation timed out after 120s")
# ── main ────────────────────────────────────────────────────────────────
async def amain():
@@ -372,6 +533,9 @@ async def amain():
cfg = AgentConfig.from_ini()
log.info("node_name=%s node_ip=%s", cfg.node_name, cfg.node_ip)
cfg.capabilities = detect_capabilities(cfg)
log.info("capabilities: %s", cfg.capabilities)
inventory = discover_models(cfg.models_dir)
log.info("discovered %d models", len(inventory))
@@ -418,6 +582,13 @@ async def amain():
await swap_queue.bind(admin_ex, f"node.{cfg.node_name}.cmd.swap_model")
await swap_queue.consume(lambda msg: handle_swap_model(cfg, channel, (admin_ex, system_ex), msg))
# Set up image gen consumer
if "image_gen" in cfg.capabilities:
image_queue = await channel.declare_queue("", exclusive=True)
await image_queue.bind(admin_ex, f"node.{cfg.node_name}.cmd.image_generate")
await image_queue.consume(lambda msg: handle_image_generate(cfg, channel, (admin_ex, system_ex), msg))
log.info("image generation handler registered")
log.info("listening for pings and commands")
# Run forever
await asyncio.Event().wait()