comfy-adapter/app/main.py

"""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()