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 Settings, get_settings
from .errors import InvalidRequest, install_error_handlers
from .injector import build_workflow, new_seed, parse_size
from .schemas import (
    ImageEditJsonRequest,
    ImageGenerationRequest,
    ImageItem,
    ImagesResponse,
    ModelCard,
    ModelsResponse,
)
from .templates import load_registry, select_model

logger = logging.getLogger("comfy-adapter")

PNG_MAGIC = b"\x89PNG\r\n\x1a\n"
JPEG_MAGIC = b"\xff\xd8\xff"

# Type-correct dummy values used to dry-inject every workflow for server validation.
_DUMMY_VALUES = {
    "prompt": "x",
    "negative": "",
    "seed": 1,
    "width": 1024,
    "height": 1024,
    "image": "dummy.png",
}


def create_app(settings: Settings | None = None) -> FastAPI:
    settings = settings or get_settings()

    @asynccontextmanager
    async def lifespan(app: FastAPI):
        registry = load_registry(settings)
        client = ComfyClient(settings.comfy_base_url)
        if settings.validate_templates_on_server:
            for model in registry.values():
                try:
                    await client.validate_graph(build_workflow(model.graph, _DUMMY_VALUES))
                    logger.info("workflow validated on ComfyUI server: %s", model.id)
                except Exception as e:  # noqa: BLE001 - startup check is advisory, logged
                    logger.error("server-side validation failed for model '%s': %s", model.id, e)
        app.state.settings = settings
        app.state.comfy = client
        app.state.registry = registry
        logger.info("comfy-adapter ready (server=%s, models=%s)", settings.comfy_base_url, list(registry))
        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))
        registry = request.app.state.registry
        return ModelsResponse(
            data=[ModelCard(id=model_id, owned_by=model.owned_by) for model_id, model in registry.items()]
        )

    @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")
        model = select_model(request.app.state.registry, body.model, "generations")
        width, height = parse_size(body.size)
        graph = build_workflow(
            model.graph,
            {"prompt": body.prompt, "seed": new_seed(body.seed), "negative": "", "width": width, "height": height},
        )
        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")
        _validate_image_bytes(image_bytes)
        model = select_model(request.app.state.registry, model, "edits")
        comfy: ComfyClient = app.state.comfy
        name = await comfy.upload_image(image_bytes)
        graph = build_workflow(
            model.graph,
            {"prompt": prompt, "seed": new_seed(seed), "negative": "", "image": name},
        )
        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")
    if upload is None or not hasattr(upload, "read"):
        # Clients following OpenAI's newer image API (e.g. LiteLLM) send the file
        # under "image[]"; also accept other common field names.
        for _key in ("image[]", "image[0]", "file", "input_file"):
            _alt = form.get(_key)
            if _alt is not None and hasattr(_alt, "read"):
                upload = _alt
                break
    # 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()