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