OpenAI-compatible adapter for ComfyUI image generation and editing

app/main.py

@@ -0,0 +1,171 @@
+"""FastAPI application: OpenAI-compatible image endpoints backed by ComfyUI."""
+import base64
+import logging
+import time
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI, Request
+
+from .auth import require_key
+from .comfy_client import ComfyClient, extract_output_image
+from .config import AVAILABLE_MODELS, CREATE_MODEL, EDIT_MODEL, Settings, get_settings
+from .errors import InvalidRequest, install_error_handlers
+from .injector import build_workflow, new_seed
+from .schemas import (
+    ImageEditJsonRequest,
+    ImageGenerationRequest,
+    ImageItem,
+    ImagesResponse,
+    ModelCard,
+    ModelsResponse,
+)
+from .templates import load_all
+
+logger = logging.getLogger("comfy-adapter")
+
+PNG_MAGIC = b"\x89PNG\r\n\x1a\n"
+JPEG_MAGIC = b"\xff\xd8\xff"
+
+
+def create_app(settings: Settings | None = None) -> FastAPI:
+    settings = settings or get_settings()
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        create_graph, edit_graph = load_all(settings)
+        client = ComfyClient(settings.comfy_base_url)
+        if settings.validate_templates_on_server:
+            try:
+                await client.validate_graph(create_graph)
+                await client.validate_graph(edit_graph)
+                logger.info("workflow templates validated on ComfyUI server")
+            except Exception as e:  # noqa: BLE001 - startup check is advisory, logged
+                logger.error("server-side template validation failed: %s", e)
+        app.state.settings = settings
+        app.state.comfy = client
+        app.state.templates = {"create": create_graph, "edit": edit_graph}
+        logger.info("comfy-adapter ready (server=%s)", settings.comfy_base_url)
+        try:
+            yield
+        finally:
+            await client.aclose()
+
+    app = FastAPI(title="comfy-adapter", lifespan=lifespan)
+    install_error_handlers(app)
+
+    def _s(req: Request) -> Settings:
+        return req.app.state.settings
+
+    async def _run(graph: dict) -> ImagesResponse:
+        comfy: ComfyClient = app.state.comfy
+        cfg: Settings = app.state.settings
+        prompt_id = await comfy.submit_prompt(graph)
+        entry = await comfy.wait_for_result(prompt_id, cfg.execution_timeout, cfg.poll_interval)
+        filename, subfolder = extract_output_image(entry)
+        image_bytes = await comfy.fetch_view(filename, subfolder)
+        return ImagesResponse(
+            created=int(time.time()),
+            data=[ImageItem(b64_json=base64.b64encode(image_bytes).decode())],
+        )
+
+    @app.get("/health")
+    async def health() -> dict:
+        return {"status": "ok"}
+
+    @app.get("/v1/models")
+    async def models(request: Request) -> ModelsResponse:
+        require_key(request, _s(request))
+        return ModelsResponse(data=[ModelCard(id=m) for m in AVAILABLE_MODELS])
+
+    @app.post("/v1/images/generations")
+    async def generations(request: Request) -> ImagesResponse:
+        settings = _s(request)
+        require_key(request, settings)
+        body = _parse_json_body(await request.body(), ImageGenerationRequest)
+        if not body.prompt.strip():
+            raise InvalidRequest("prompt must not be empty")
+        if body.n not in (None, 1):
+            raise InvalidRequest("only n=1 is supported")
+        if body.model is not None and body.model != CREATE_MODEL:
+            raise InvalidRequest(f"unknown model '{body.model}', expected '{CREATE_MODEL}'")
+        graph = build_workflow(
+            app.state.templates["create"],
+            settings.mapping()["create"],
+            prompt=body.prompt,
+            seed=new_seed(body.seed),
+            size=body.size,
+        )
+        return await _run(graph)
+
+    @app.post("/v1/images/edits")
+    async def edits(request: Request) -> ImagesResponse:
+        settings = _s(request)
+        require_key(request, settings)
+        content_type = request.headers.get("content-type", "")
+        if content_type.startswith("multipart/form-data"):
+            image_bytes, prompt, model, seed = await _parse_multipart(request)
+        elif content_type.startswith("application/json"):
+            body = _parse_json_body(await request.body(), ImageEditJsonRequest)
+            image_bytes = base64.b64decode(body.image)
+            prompt, model, seed = body.prompt, body.model, body.seed
+        else:
+            raise InvalidRequest("expected multipart/form-data or application/json body")
+        if not prompt.strip():
+            raise InvalidRequest("prompt must not be empty")
+        if model is not None and model != EDIT_MODEL:
+            raise InvalidRequest(f"unknown model '{model}' for edits, expected '{EDIT_MODEL}'")
+        _validate_image_bytes(image_bytes)
+        comfy: ComfyClient = app.state.comfy
+        name = await comfy.upload_image(image_bytes)
+        graph = build_workflow(
+            app.state.templates["edit"],
+            settings.mapping()["edit"],
+            prompt=prompt,
+            seed=new_seed(seed),
+            image_name=name,
+            negative="",
+        )
+        return await _run(graph)
+
+    return app
+
+
+def _parse_json_body(raw: bytes, model_cls):
+    from .errors import InvalidRequest
+    import json
+
+    try:
+        data = json.loads(raw or b"{}")
+    except json.JSONDecodeError as e:
+        raise InvalidRequest(f"invalid JSON body: {e}") from e
+    try:
+        return model_cls.model_validate(data)
+    except Exception as e:
+        raise InvalidRequest(f"invalid request body: {e}") from e
+
+
+async def _parse_multipart(request: Request):
+    form = await request.form()
+    upload = form.get("image")
+    # duck-typing: FastAPI 0.142+ ships its own UploadFile class, forms yield starlette's
+    if upload is None or not hasattr(upload, "read"):
+        raise InvalidRequest("field 'image' (file) is required")
+    image_bytes = await upload.read()
+    prompt = str(form.get("prompt") or "")
+    model = form.get("model")
+    seed_raw = form.get("seed")
+    try:
+        seed = int(seed_raw) if seed_raw not in (None, "") else None
+    except (TypeError, ValueError):
+        raise InvalidRequest("seed must be an integer")
+    return image_bytes, prompt, (str(model) if model else None), seed
+
+
+def _validate_image_bytes(data: bytes) -> None:
+    if not data:
+        raise InvalidRequest("image is empty")
+    if not (data.startswith(PNG_MAGIC) or data.startswith(JPEG_MAGIC)):
+        raise InvalidRequest("image must be PNG or JPEG")
+
+
+app = create_app()