OpenAI-compatible adapter for ComfyUI image generation and editing

app/auth.py

@@ -0,0 +1,15 @@
+"""Static bearer-key authentication."""
+import secrets
+
+from fastapi import Request
+
+from .config import Settings
+from .errors import Unauthorized
+
+
+def require_key(request: Request, settings: Settings) -> None:
+    header = request.headers.get("authorization", "")
+    scheme, _, token = header.partition(" ")
+    expected = settings.adapter_api_key
+    if not expected or scheme.lower() != "bearer" or not secrets.compare_digest(token, expected):
+        raise Unauthorized()

app/comfy_client.py

@@ -0,0 +1,124 @@
+"""Async HTTP client for the ComfyUI server (official endpoints only)."""
+import asyncio
+import json
+import secrets
+import time
+import uuid
+from typing import Any
+
+import httpx
+
+from .errors import ComfyUIError
+
+
+class ComfyClient:
+    def __init__(self, base_url: str, timeout: float = 30.0):
+        self._client = httpx.AsyncClient(base_url=base_url, timeout=timeout)
+
+    async def aclose(self) -> None:
+        await self._client.aclose()
+
+    async def submit_prompt(self, graph: dict[str, Any]) -> str:
+        """Queue a graph; returns prompt_id. ComfyUI validates synchronously."""
+        payload = {"prompt": graph, "client_id": str(uuid.uuid4())}
+        try:
+            resp = await self._client.post("/prompt", json=payload)
+        except httpx.HTTPError as e:
+            raise ComfyUIError(f"ComfyUI unreachable: {e}") from e
+        if resp.status_code == 400:
+            detail = _validation_detail(resp.json())
+            raise ComfyUIError(f"ComfyUI rejected workflow: {detail}", code="comfyui_validation_failed")
+        if resp.status_code != 200:
+            raise ComfyUIError(f"ComfyUI /prompt returned HTTP {resp.status_code}")
+        prompt_id = resp.json().get("prompt_id")
+        if not prompt_id:
+            raise ComfyUIError("ComfyUI /prompt response missing prompt_id")
+        return prompt_id
+
+    async def upload_image(self, data: bytes, content_type: str = "image/png") -> str:
+        """Upload an input image; returns the server-side filename."""
+        filename = f"adapter_{uuid.uuid4().hex[:8]}_{secrets.token_hex(4)}.png"
+        try:
+            resp = await self._client.post(
+                "/upload/image",
+                files={"image": (filename, data, content_type)},
+            )
+        except httpx.HTTPError as e:
+            raise ComfyUIError(f"ComfyUI unreachable: {e}") from e
+        if resp.status_code != 200:
+            raise ComfyUIError(f"ComfyUI /upload/image returned HTTP {resp.status_code}")
+        name = resp.json().get("name")
+        if not name:
+            raise ComfyUIError("ComfyUI /upload/image response missing name")
+        return name
+
+    async def wait_for_result(self, prompt_id: str, timeout: float, interval: float) -> dict:
+        """Poll /history/{prompt_id} until the entry appears; returns the history entry."""
+        deadline = time.monotonic() + timeout
+        while True:
+            try:
+                resp = await self._client.get(f"/history/{prompt_id}")
+            except httpx.HTTPError as e:
+                raise ComfyUIError(f"ComfyUI unreachable: {e}") from e
+            if resp.status_code != 200:
+                raise ComfyUIError(f"ComfyUI /history returned HTTP {resp.status_code}")
+            entry = resp.json().get(prompt_id)
+            if entry is not None:
+                status = entry.get("status", {})
+                if status.get("status_str") == "error" or status.get("status_str") == "execution_failed":
+                    raise ComfyUIError(f"ComfyUI execution failed for prompt {prompt_id}", code="comfyui_execution_failed")
+                return entry
+            if time.monotonic() >= deadline:
+                raise ComfyUIError(
+                    f"Timed out after {timeout:.0f}s waiting for prompt {prompt_id}",
+                    code="comfyui_timeout",
+                )
+            await asyncio.sleep(interval)
+
+    async def fetch_view(self, filename: str, subfolder: str = "") -> bytes:
+        """Download a produced image from the output folder."""
+        try:
+            resp = await self._client.get(
+                "/view", params={"filename": filename, "subfolder": subfolder, "type": "output"}
+            )
+        except httpx.HTTPError as e:
+            raise ComfyUIError(f"ComfyUI unreachable: {e}") from e
+        if resp.status_code != 200:
+            raise ComfyUIError(f"ComfyUI /view returned HTTP {resp.status_code} for '{filename}'")
+        return resp.content
+
+    async def validate_graph(self, graph: dict[str, Any]) -> None:
+        """Validate a graph without queueing (POST /prompt?front=false)."""
+        try:
+            resp = await self._client.post("/prompt", params={"front": "false"}, json={"prompt": graph})
+        except httpx.HTTPError as e:
+            raise ComfyUIError(f"ComfyUI unreachable: {e}") from e
+        if resp.status_code == 200:
+            return
+        if resp.status_code == 400:
+            raise ComfyUIError(
+                f"ComfyUI rejected workflow: {_validation_detail(resp.json())}",
+                code="comfyui_validation_failed",
+            )
+        raise ComfyUIError(f"ComfyUI /prompt?front=false returned HTTP {resp.status_code}")
+
+
+def extract_output_image(history_entry: dict) -> tuple[str, str]:
+    """(filename, subfolder) of the first output image, preferring SaveImage nodes."""
+    outputs = history_entry.get("outputs", {})
+    images: list[dict] = []
+    for node_output in outputs.values():
+        images.extend(img for img in node_output.get("images", []) if img.get("type") == "output")
+    if not images:
+        raise ComfyUIError("ComfyUI produced no output images", code="comfyui_no_output")
+    return images[0]["filename"], images[0].get("subfolder", "")
+
+
+def _validation_detail(body: Any) -> str:
+    if isinstance(body, dict) and "node_errors" in body:
+        parts = []
+        for node_id, node_err in body["node_errors"].items():
+            for err in node_err.get("errors", []):
+                parts.append(f"node {node_id}: {err.get('message')} {err.get('details', '')}".strip())
+        return "; ".join(parts) or "validation failed"
+    return json.dumps(body)[:300] if body else "validation failed"

app/config.py

@@ -0,0 +1,55 @@
+"""Environment settings and the parameter -> workflow node mapping."""
+import json
+from pathlib import Path
+
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+
+CREATE_MODEL = "z-image-turbo"
+EDIT_MODEL = "qwen-image-edit"
+AVAILABLE_MODELS = [CREATE_MODEL, EDIT_MODEL]
+
+# Node ids refer to workflow_create.json / workflow_edit.json.
+# Override the whole map via MAPPING_FILE (JSON with "create" and "edit" keys).
+DEFAULT_MAPPING: dict[str, dict[str, dict[str, str]]] = {
+    "create": {
+        "prompt": {"node": "83:27", "field": "text"},
+        "size": {"node": "83:13", "width_field": "width", "height_field": "height"},
+        "seed": {"node": "83:3", "field": "seed"},
+    },
+    "edit": {
+        "prompt": {"node": "115:111", "field": "prompt"},
+        "negative": {"node": "115:110", "field": "prompt"},
+        "seed": {"node": "115:3", "field": "seed"},
+        "image": {"node": "78", "field": "image"},
+    },
+}
+
+
+class Settings(BaseSettings):
+    model_config = SettingsConfigDict(env_file=".env", extra="ignore")
+
+    comfy_base_url: str = "http://192.168.0.3:8188"
+    adapter_api_key: str = ""
+    execution_timeout: float = 300.0
+    poll_interval: float = 1.5
+    create_template_path: Path = REPO_ROOT / "workflow_create.json"
+    edit_template_path: Path = REPO_ROOT / "workflow_edit.json"
+    mapping_file: Path | None = None
+    # Validate graphs on the ComfyUI server (POST /prompt?front=false) at startup.
+    validate_templates_on_server: bool = True
+
+    def mapping(self) -> dict[str, dict[str, dict[str, str]]]:
+        if self.mapping_file is None:
+            return DEFAULT_MAPPING
+        with open(self.mapping_file, encoding="utf-8") as f:
+            loaded = json.load(f)
+        merged = {k: dict(v) for k, v in DEFAULT_MAPPING.items()}
+        for section, section_map in loaded.items():
+            merged.setdefault(section, {}).update(section_map)
+        return merged
+
+
+def get_settings() -> Settings:
+    return Settings()

app/errors.py

@@ -0,0 +1,41 @@
+"""Adapter exceptions and OpenAI-shaped error responses."""
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+
+
+class AdapterError(Exception):
+    """Base class; renders as an OpenAI error body."""
+
+    def __init__(self, status: int, message: str, err_type: str, code: str | None = None):
+        self.status = status
+        self.message = message
+        self.err_type = err_type
+        self.code = code
+        super().__init__(message)
+
+
+class InvalidRequest(AdapterError):
+    def __init__(self, message: str, code: str | None = None):
+        super().__init__(400, message, "invalid_request_error", code)
+
+
+class Unauthorized(AdapterError):
+    def __init__(self, message: str = "Invalid API key provided"):
+        super().__init__(401, message, "invalid_request_error", "invalid_api_key")
+
+
+class ComfyUIError(AdapterError):
+    """ComfyUI rejected validation, failed execution, or timed out."""
+
+    def __init__(self, message: str, code: str = "comfyui_error"):
+        super().__init__(502, message, "server_error", code)
+
+
+def error_body(message: str, err_type: str, code: str | None = None) -> dict:
+    return {"error": {"message": message, "type": err_type, "param": None, "code": code}}
+
+
+def install_error_handlers(app: FastAPI) -> None:
+    @app.exception_handler(AdapterError)
+    async def adapter_error_handler(request: Request, exc: AdapterError) -> JSONResponse:
+        return JSONResponse(status_code=exc.status, content=error_body(exc.message, exc.err_type, exc.code))

app/injector.py

@@ -0,0 +1,67 @@
+"""Build a concrete workflow graph from a template + request parameters."""
+import copy
+import secrets
+from typing import Any
+
+from .errors import InvalidRequest
+
+SEED_MAX = 2**63 - 1
+
+
+def parse_size(size: str) -> tuple[int, int]:
+    """Parse OpenAI `WxH` size strings (e.g. '1024x1024')."""
+    try:
+        w_str, h_str = size.lower().split("x", 1)
+        width, height = int(w_str), int(h_str)
+    except ValueError as e:
+        raise InvalidRequest(f"invalid size '{size}', expected format 'WIDTHxHEIGHT'") from e
+    if width <= 0 or height <= 0:
+        raise InvalidRequest(f"invalid size '{size}', dimensions must be positive")
+    return width, height
+
+
+def new_seed(seed: int | None) -> int:
+    if seed is None:
+        return secrets.randbelow(SEED_MAX) + 1
+    if not 0 <= seed <= SEED_MAX:
+        raise InvalidRequest(f"seed must be between 0 and {SEED_MAX}")
+    return seed
+
+
+def build_workflow(
+    template: dict[str, Any],
+    mapping: dict[str, dict[str, str]],
+    *,
+    prompt: str,
+    seed: int,
+    size: str | None = None,
+    image_name: str | None = None,
+    negative: str = "",
+) -> dict[str, Any]:
+    """Deep-copy the template and inject request values per the node mapping."""
+    graph = copy.deepcopy(template)
+
+    def put(key: str, value: Any) -> None:
+        spec = mapping.get(key)
+        if spec is None:
+            return
+        node = graph.get(spec["node"])
+        if node is None:
+            raise InvalidRequest(f"mapping key '{key}' points to missing node '{spec['node']}'")
+        node["inputs"][spec["field"]] = value
+
+    put("prompt", prompt)
+    put("negative", negative)
+    put("seed", seed)
+    if image_name is not None:
+        put("image", image_name)
+    if size is not None:
+        spec = mapping.get("size")
+        if spec is not None:
+            width, height = parse_size(size)
+            node = graph.get(spec["node"])
+            if node is None:
+                raise InvalidRequest(f"mapping key 'size' points to missing node '{spec['node']}'")
+            node["inputs"][spec["width_field"]] = width
+            node["inputs"][spec["height_field"]] = height
+    return graph

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

app/schemas.py

@@ -0,0 +1,62 @@
+"""Pydantic models for the OpenAI images contract."""
+import base64
+import binascii
+
+from pydantic import BaseModel, field_validator
+
+DEFAULT_SIZE = "1024x1024"
+
+
+class ImageGenerationRequest(BaseModel):
+    prompt: str
+    model: str | None = None
+    n: int | None = 1
+    size: str = DEFAULT_SIZE
+    seed: int | None = None
+    response_format: str | None = None  # accepted; only b64_json is produced
+
+    @field_validator("size")
+    @classmethod
+    def _size_shape(cls, v: str) -> str:
+        w_str, sep, h_str = v.lower().partition("x")
+        if not sep or not w_str.isdigit() or not h_str.isdigit() or int(w_str) <= 0 or int(h_str) <= 0:
+            raise ValueError(f"invalid size '{v}', expected format 'WIDTHxHEIGHT'")
+        return v
+
+
+class ImageEditJsonRequest(BaseModel):
+    image: str  # base64-encoded PNG/JPEG (data-URL prefix allowed)
+    prompt: str
+    model: str | None = None
+    size: str | None = None  # accepted but ignored: edits size comes from the input image
+    seed: int | None = None
+
+    @field_validator("image")
+    @classmethod
+    def _decode_b64(cls, v: str) -> str:
+        raw = v.split(",", 1)[1] if v.startswith("data:") else v
+        try:
+            base64.b64decode(raw, validate=True)
+        except (binascii.Error, ValueError) as e:
+            raise ValueError("image must be valid base64") from e
+        return raw
+
+
+class ImageItem(BaseModel):
+    b64_json: str
+
+
+class ImagesResponse(BaseModel):
+    created: int
+    data: list[ImageItem]
+
+
+class ModelCard(BaseModel):
+    id: str
+    object: str = "model"
+    owned_by: str = "comfy-adapter"
+
+
+class ModelsResponse(BaseModel):
+    object: str = "list"
+    data: list[ModelCard]

app/templates.py

@@ -0,0 +1,62 @@
+"""Template loading with startup validation of graph links and mappings."""
+import json
+from pathlib import Path
+
+from .config import Settings
+from .errors import AdapterError
+
+
+class TemplateError(AdapterError):
+    """Malformed workflow template; blocks startup."""
+
+    def __init__(self, message: str):
+        super().__init__(500, message, "server_error", "template_error")
+
+
+def find_broken_links(graph: dict) -> list[str]:
+    """References to node ids that do not exist in the graph."""
+    errors = []
+    for node_id, node in graph.items():
+        for field, value in node.get("inputs", {}).items():
+            if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str):
+                if value[0] not in graph:
+                    errors.append(f"node {node_id}: input '{field}' links to missing node '{value[0]}'")
+    return errors
+
+
+def check_mapping(graph: dict, mapping: dict, name: str) -> list[str]:
+    """Every mapped node id must exist with a scalar input field."""
+    errors = []
+    for key, spec in mapping.items():
+        node = graph.get(spec["node"])
+        if node is None:
+            errors.append(f"mapping '{name}.{key}': node '{spec['node']}' not found in template")
+            continue
+        for field in _mapped_fields(spec):
+            if field not in node.get("inputs", {}) or isinstance(node["inputs"][field], list):
+                errors.append(f"mapping '{name}.{key}': node '{spec['node']}' has no scalar input '{field}'")
+    return errors
+
+
+def _mapped_fields(spec: dict) -> list[str]:
+    if "width_field" in spec:
+        return [spec["width_field"], spec["height_field"]]
+    return [spec["field"]]
+
+
+def load_template(path: Path, graph_name: str, mapping: dict, mapping_name: str) -> dict:
+    with open(path, encoding="utf-8") as f:
+        graph = json.load(f)
+    errors = find_broken_links(graph) + check_mapping(graph, mapping[mapping_name], mapping_name)
+    if errors:
+        raise TemplateError(f"workflow template '{graph_name}' ({path}) is invalid:\n  " + "\n  ".join(errors))
+    return graph
+
+
+def load_all(settings: Settings) -> tuple[dict, dict]:
+    mapping = settings.mapping()
+    create = load_template(
+        settings.create_template_path, "workflow_create.json", mapping, "create"
+    )
+    edit = load_template(settings.edit_template_path, "workflow_edit.json", mapping, "edit")
+    return create, edit