71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
"""cAIc 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}
|