From c17ab5036ecf2ed90e92bc28e87d44db3c9c0eb8 Mon Sep 17 00:00:00 2001 From: Shaman Date: Sun, 4 Oct 2026 21:56:00 +0300 Subject: [PATCH] OpenAI-compatible adapter for ComfyUI image generation and editing --- .dockerignore | 7 ++ .env.example | 10 +++ .gitignore | 7 ++ Dockerfile | 12 +++ README.md | 134 +++++++++++++++++++++++++++++ app/__init__.py | 0 app/auth.py | 15 ++++ app/comfy_client.py | 124 +++++++++++++++++++++++++++ app/config.py | 55 ++++++++++++ app/errors.py | 41 +++++++++ app/injector.py | 67 +++++++++++++++ app/main.py | 171 +++++++++++++++++++++++++++++++++++++ app/schemas.py | 62 ++++++++++++++ app/templates.py | 62 ++++++++++++++ pyproject.toml | 31 +++++++ tests/__init__.py | 0 tests/conftest.py | 51 +++++++++++ tests/test_auth.py | 26 ++++++ tests/test_comfy_client.py | 120 ++++++++++++++++++++++++++ tests/test_config.py | 32 +++++++ tests/test_edits.py | 107 +++++++++++++++++++++++ tests/test_errors.py | 29 +++++++ tests/test_generations.py | 140 ++++++++++++++++++++++++++++++ tests/test_injector.py | 73 ++++++++++++++++ tests/test_live.py | 70 +++++++++++++++ tests/test_models.py | 11 +++ tests/test_templates.py | 57 +++++++++++++ workflow_create.json | 78 +++++++++++++++++ workflow_edit.json | 138 ++++++++++++++++++++++++++++++ 29 files changed, 1730 insertions(+) create mode 100644 .dockerignore create mode 100644 .env.example create mode 100644 .gitignore create mode 100644 Dockerfile create mode 100644 README.md create mode 100644 app/__init__.py create mode 100644 app/auth.py create mode 100644 app/comfy_client.py create mode 100644 app/config.py create mode 100644 app/errors.py create mode 100644 app/injector.py create mode 100644 app/main.py create mode 100644 app/schemas.py create mode 100644 app/templates.py create mode 100644 pyproject.toml create mode 100644 tests/__init__.py create mode 100644 tests/conftest.py create mode 100644 tests/test_auth.py create mode 100644 tests/test_comfy_client.py create mode 100644 tests/test_config.py create mode 100644 tests/test_edits.py create mode 100644 tests/test_errors.py create mode 100644 tests/test_generations.py create mode 100644 tests/test_injector.py create mode 100644 tests/test_live.py create mode 100644 tests/test_models.py create mode 100644 tests/test_templates.py create mode 100644 workflow_create.json create mode 100644 workflow_edit.json diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..1482899 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,7 @@ +.venv +__pycache__ +*.pyc +.pytest_cache +.env +tests +*.md diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..2217c2e --- /dev/null +++ b/.env.example @@ -0,0 +1,10 @@ +# ComfyUI server (required) +COMFY_BASE_URL=http://192.168.0.3:8188 +# Static bearer key clients must send (required) +ADAPTER_API_KEY=sk-1234 +# Seconds to wait for ComfyUI execution result +EXECUTION_TIMEOUT=300 +# Poll interval for /history, seconds +POLL_INTERVAL=1.5 +# Optional: JSON file overriding node mappings (see app/config.py defaults) +# MAPPING_FILE=mapping.json diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b883b87 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +.venv/ +__pycache__/ +*.pyc +.env +.pytest_cache/ +*.egg-info/ +dist/ diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..2eee778 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,12 @@ +FROM python:3.12-slim + +WORKDIR /app +ENV PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1 + +COPY pyproject.toml ./ +COPY app ./app +COPY workflow_create.json workflow_edit.json ./ +RUN pip install --no-cache-dir . + +EXPOSE 8000 +CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"] diff --git a/README.md b/README.md new file mode 100644 index 0000000..ddc9201 --- /dev/null +++ b/README.md @@ -0,0 +1,134 @@ +# comfy-adapter + +OpenAI-совместимый адаптер поверх ComfyUI: `POST /v1/images/generations` (генерация, +модель `z-image-turbo`) и `POST /v1/images/edits` (редактирование, модель +`qwen-image-edit`). Адаптер подставляет параметры запроса в шаблоны workflow, +отправляет их в ComfyUI, дождаться выполнения и возвращает изображение как +`data[0].b64_json` (base64 PNG). + +``` +Client --(OpenAI contract)--> comfy-adapter --(ComfyUI HTTP API)--> ComfyUI + /v1/images/generations POST /prompt + /v1/images/edits POST /upload/image + GET /history/{id} + GET /view +``` + +## Запуск + +```bash +python -m venv .venv && .venv/bin/pip install -e . # Windows: .venv\Scripts\ +cp .env.example .env # заполните COMFY_BASE_URL, ADAPTER_API_KEY +uvicorn app.main:app --host 0.0.0.0 --port 8000 +``` + +Или Docker: + +```bash +docker build -t comfy-adapter . +docker run -p 8000:8000 --env-file .env comfy-adapter +``` + +При старте адаптер загружает `workflow_create.json` / `workflow_edit.json`, локально +проверяет графы на битые ссылки (битый шаблон — падение старта со списком проблемных +узлов) и прогоняет серверную проверку `POST /prompt?front=false` (отключается: +`VALIDATE_TEMPLATES_ON_SERVER=false`). Логи стартовой проверки видны с +`uvicorn --log-level info`. + +## Переменные окружения + +| Переменная | Default | Назначение | +|---|---|---| +| `COMFY_BASE_URL` | `http://192.168.0.3:8188` | Адрес ComfyUI | +| `ADAPTER_API_KEY` | — (обязательна) | Статический Bearer-ключ клиентов | +| `EXECUTION_TIMEOUT` | `300` | Секунд ожидания результата ComfyUI | +| `POLL_INTERVAL` | `1.5` | Интервал опроса `/history`, сек | +| `MAPPING_FILE` | — | JSON-файл, переопределяющий карту маппинга | +| `VALIDATE_TEMPLATES_ON_SERVER` | `true` | Проверять шаблоны на сервере при старте | + +## Контракт + +`POST /v1/images/generations` (JSON, `Authorization: Bearer `): + +```json +{"prompt": "a red cube", "model": "z-image-turbo", "n": 1, "size": "1024x1024", "seed": 123} +``` + +- `n` поддерживается только 1 (иначе 400); `seed` необязателен — без него случайный + (повтор с тем же seed + prompt даёт идентичное изображение). +- Ответ: `{"created": , "data": [{"b64_json": ""}]}`. + +`POST /v1/images/edits` — два транспорта: + +- multipart: поля `image` (файл PNG/JPEG) + `prompt` (+ необязательные `model`, `seed`); +- JSON: `{"image": "", "prompt": "...", "model": "qwen-image-edit", "seed": 5}`. + +Поле `size` в edits принимается, но не влияет на выход: итоговый размер — ~1 Мп с +аспектом входного изображения (это делает сам workflow). + +Ошибки — в формате OpenAI (`error.message/type/param/code`): `400` — невалидное тело; +`401` — ключ отсутствует/неверен; `502` — ComfyUI недоступен, отверг workflow +(`comfyui_validation_failed`), ошибко исполнения (`comfyui_execution_failed`) или таймаут +(`comfyui_timeout`). `GET /health` — без ключа. + +```bash +# генерация +curl -s http://localhost:8000/v1/images/generations \ + -H "Authorization: Bearer $ADAPTER_API_KEY" -H "Content-Type: application/json" \ + -d '{"prompt":"a red cube on white background","size":"1024x1024","seed":42}' | jq -r .data[0].b64_json | base64 -d > out.png + +# редактирование +curl -s http://localhost:8000/v1/images/edits \ + -H "Authorization: Bearer $ADAPTER_API_KEY" \ + -F image=@input.png -F prompt="make the background deep blue" | jq -r .data[0].b64_json | base64 -d > edited.png +``` + +```python +from openai import OpenAI +client = OpenAI(base_url="http://localhost:8000/v1", api_key="sk-...") +r = client.images.generate(model="z-image-turbo", prompt="a red cube", size="1024x1024", n=1) +open("out.png", "wb").write(__import__("base64").b64decode(r.data[0].b64_json)) +``` + +## Карта маппинга + +Параметр -> узел workflow (правится без кода через `MAPPING_FILE`): + +| Параметр | workflow_create.json | workflow_edit.json | +|---|---|---| +| prompt | `83:27.inputs.text` | `115:111.inputs.prompt` | +| negative | — (ConditioningZeroOut) | `115:110.inputs.prompt` (всегда "") | +| size | `83:13.inputs.width/height` | — (размер = ~1 Мп входа) | +| seed | `83:3.inputs.seed` | `115:3.inputs.seed` | +| image | — | upload -> `78.inputs.image` | + +## Починка шаблонов относительно экспорта OpenWebUI + +Выгруженные из OpenWebUI шаблоны **не проходили валидацию ComfyUI** (`KeyError` по +несуществующим узлам). Исправления в репозиторных файлах: + +- `workflow_create.json`: ссылки `[83:43, 0]`/`[83:43, 1]` (узла нет в графе) + переподключены на существующие лоадеры: `83:3.inputs.model -> [83:28, 0]` + (UNETLoader z_image_turbo), `83:27.inputs.clip -> [83:30, 0]` (CLIPLoader qwen_3_4b). +- `workflow_edit.json`: `115:66.inputs.model -> [115:89, 0]` вместо несуществующего + `115:117` — LoRA `Qwen-Image-Edit-2509-Lightning-4steps` (нода `115:89`, в экспорте + «висячая») подключена в цепочку модели. + +Оба шаблона проходят `POST /prompt?front=false` и реально генерируют (проверено на +ComfyUI 0.37.0 / V100). + +## Замечания по продакшену + +- Ответ синхронный: соединение висит до конца генерации. За reverse proxy поднимите + `proxy_read_timeout` (nginx) / эквивалент выше, чем `EXECUTION_TIMEOUT`. +- Фоновые интеграции (OpenWebUI) продолжают работать с ComfyUI напрямую — адаптер + ничего не migрирует. + +## Тесты + +```bash +pip install -e ".[dev]" +pytest # unit + mocked-ComfyUI (respx) +pytest -m live # живые прогоны: нужен запущенный адаптер и ADAPTER_LIVE=1, + # ADAPTER_URL=http://127.0.0.1:8000 ADAPTER_KEY=<ключ> +``` diff --git a/app/__init__.py b/app/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/app/auth.py b/app/auth.py new file mode 100644 index 0000000..24ffe4a --- /dev/null +++ b/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() diff --git a/app/comfy_client.py b/app/comfy_client.py new file mode 100644 index 0000000..152ca12 --- /dev/null +++ b/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" diff --git a/app/config.py b/app/config.py new file mode 100644 index 0000000..dc5d37e --- /dev/null +++ b/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() diff --git a/app/errors.py b/app/errors.py new file mode 100644 index 0000000..e8394c8 --- /dev/null +++ b/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)) diff --git a/app/injector.py b/app/injector.py new file mode 100644 index 0000000..c543aae --- /dev/null +++ b/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 diff --git a/app/main.py b/app/main.py new file mode 100644 index 0000000..8af1be7 --- /dev/null +++ b/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() diff --git a/app/schemas.py b/app/schemas.py new file mode 100644 index 0000000..42b5256 --- /dev/null +++ b/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] diff --git a/app/templates.py b/app/templates.py new file mode 100644 index 0000000..332bd0e --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..1aa6951 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,31 @@ +[project] +name = "comfy-adapter" +version = "0.1.0" +description = "OpenAI-compatible adapter for ComfyUI image generation and editing" +requires-python = ">=3.12" +dependencies = [ + "fastapi>=0.115", + "uvicorn[standard]>=0.30", + "httpx>=0.27", + "pydantic>=2.7", + "pydantic-settings>=2.3", + "python-multipart>=0.0.9", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8", + "pytest-asyncio>=0.23", + "respx>=0.21", +] + +[tool.pytest.ini_options] +testpaths = ["tests"] +asyncio_mode = "auto" +asyncio_default_fixture_loop_scope = "function" +markers = [ + "live: tests against the real ComfyUI server (manual run)", +] + +[tool.setuptools] +packages = ["app"] diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..321faed --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,51 @@ +import struct +import zlib + +import httpx +import pytest + +from app.config import Settings +from app.main import create_app + +COMFY = "http://comfy.test" +FAKE_KEY = "sk-test-1234" + + +def settings(**overrides) -> Settings: + base = dict( + comfy_base_url=COMFY, + adapter_api_key=FAKE_KEY, + validate_templates_on_server=False, + poll_interval=0.0, + execution_timeout=5.0, + ) + base.update(overrides) + return Settings(**base) + + +def png_bytes(w: int = 64, h: int = 64) -> bytes: + def chunk(t: bytes, d: bytes) -> bytes: + return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d)) + + ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0) + iend = b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", ihdr) + chunk(b"IDAT", zlib.compress(b"\x00" * (w * 3 + 1) * h)) + chunk(b"IEND", b"") + return iend + + +def png_size(data: bytes) -> tuple[int, int]: + return struct.unpack(">II", data[16:24]) + + +@pytest.fixture +async def client(): + """Async client bound to the app with lifespan running + (templates loaded, comfy client targeting http://comfy.test).""" + import contextlib + + app = create_app(settings()) + async with contextlib.AsyncExitStack() as stack: + await stack.enter_async_context(app.router.lifespan_context(app)) + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://adapter.test") as ac: + ac.headers["Authorization"] = f"Bearer {FAKE_KEY}" + yield ac diff --git a/tests/test_auth.py b/tests/test_auth.py new file mode 100644 index 0000000..a1b1e7b --- /dev/null +++ b/tests/test_auth.py @@ -0,0 +1,26 @@ +import httpx +import pytest + +from tests.conftest import FAKE_KEY + + +async def test_missing_key_401(client: httpx.AsyncClient): + r = await client.post("/v1/images/generations", json={"prompt": "x"}, headers={"Authorization": ""}) + assert r.status_code == 401 + assert r.json()["error"]["code"] == "invalid_api_key" + + +async def test_wrong_key_401(client: httpx.AsyncClient): + r = await client.post("/v1/images/generations", json={"prompt": "x"}, headers={"Authorization": "Bearer sk-wrong"}) + assert r.status_code == 401 + + +async def test_valid_key_accepted(client: httpx.AsyncClient): + r = await client.get("/v1/models") + assert r.status_code == 200 + + +async def test_health_open_without_key(client: httpx.AsyncClient): + r = await client.get("/health", headers={"Authorization": ""}) + assert r.status_code == 200 + assert r.json() == {"status": "ok"} diff --git a/tests/test_comfy_client.py b/tests/test_comfy_client.py new file mode 100644 index 0000000..3d05c66 --- /dev/null +++ b/tests/test_comfy_client.py @@ -0,0 +1,120 @@ +import httpx +import pytest +import respx + +from app.comfy_client import ComfyClient, extract_output_image +from app.errors import ComfyUIError + +BASE = "http://comfy.test" + + +@pytest.fixture +def comfy(): + client = ComfyClient(BASE) + yield client + + +async def _history_entry(filename="out_001.png", subfolder=""): + return {"outputs": {"60": {"images": [{"filename": filename, "subfolder": subfolder, "type": "output"}]}}, + "status": {"status_str": "success", "completed": True}} + + +@respx.mock +async def test_submit_prompt_ok(comfy): + route = respx.post(f"{BASE}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": "p-1"})) + assert await comfy.submit_prompt({"1": {"class_type": "X", "inputs": {}}}) == "p-1" + body = route.calls.last.request.content.decode() + assert '"client_id"' in body + + +@respx.mock +async def test_submit_prompt_validation_error(comfy): + respx.post(f"{BASE}/prompt").mock( + return_value=httpx.Response(400, json={"error": {}, "node_errors": {"3": {"errors": [{"message": "bad link", "details": "'83:43'"}]}}}) + ) + with pytest.raises(ComfyUIError, match="83:43"): + await comfy.submit_prompt({}) + + +@respx.mock +async def test_submit_prompt_server_error(comfy): + respx.post(f"{BASE}/prompt").mock(return_value=httpx.Response(500)) + with pytest.raises(ComfyUIError, match="HTTP 500"): + await comfy.submit_prompt({}) + + +@respx.mock +async def test_upload_image_ok_returns_name(comfy): + route = respx.post(f"{BASE}/upload/image").mock(return_value=httpx.Response(200, json={"name": "adapter_x.png", "type": "input"})) + name = await comfy.upload_image(b"pngbytes") + assert name == "adapter_x.png" + assert "multipart" in route.calls.last.request.headers["content-type"] + assert b"adapter_" in route.calls.last.request.content + + +@respx.mock +async def test_upload_image_error(comfy): + respx.post(f"{BASE}/upload/image").mock(return_value=httpx.Response(500)) + with pytest.raises(ComfyUIError): + await comfy.upload_image(b"x") + + +@respx.mock +async def test_wait_for_result_returns_entry(comfy): + respx.get(f"{BASE}/history/p-1").mock(return_value=httpx.Response(200, json={"p-1": await _history_entry()})) + entry = await comfy.wait_for_result("p-1", timeout=5, interval=0) + assert entry["status"]["completed"] is True + + +@respx.mock +async def test_wait_for_result_execution_failed(comfy): + respx.get(f"{BASE}/history/p-1").mock( + return_value=httpx.Response(200, json={"p-1": {"outputs": {}, "status": {"status_str": "error", "completed": False}}}) + ) + with pytest.raises(ComfyUIError, match="execution failed"): + await comfy.wait_for_result("p-1", timeout=5, interval=0) + + +@respx.mock +async def test_wait_for_result_timeout(comfy): + respx.get(f"{BASE}/history/p-1").mock(return_value=httpx.Response(200, json={})) + with pytest.raises(ComfyUIError, match="Timed out"): + await comfy.wait_for_result("p-1", timeout=0.1, interval=0.01) + + +@respx.mock +async def test_fetch_view_ok(comfy): + respx.get(f"{BASE}/view").mock(return_value=httpx.Response(200, content=b"PNGDATA")) + assert await comfy.fetch_view("f.png") == b"PNGDATA" + + +@respx.mock +async def test_fetch_view_error(comfy): + respx.get(f"{BASE}/view").mock(return_value=httpx.Response(404)) + with pytest.raises(ComfyUIError): + await comfy.fetch_view("missing.png") + + +@respx.mock +async def test_validate_graph_ok_uses_front_false(comfy): + route = respx.post(f"{BASE}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": "x"})) + await comfy.validate_graph({"1": {"class_type": "X", "inputs": {}}}) + assert "front=false" in route.calls.last.request.url.query.decode() or "front" in str(route.calls.last.request.url) + + +@respx.mock +async def test_validate_graph_rejects(comfy): + respx.post(f"{BASE}/prompt").mock(return_value=httpx.Response(400, json={"node_errors": {"2": {"errors": [{"message": "missing", "details": ""}]}}})) + with pytest.raises(ComfyUIError, match="missing"): + await comfy.validate_graph({}) + + +def test_extract_output_image_prefers_output_type(): + entry = {"outputs": {"77": {"images": [{"filename": "prev.png", "type": "temp"}]}, + "60": {"images": [{"filename": "real.png", "subfolder": "s", "type": "output"}]}}} + assert extract_output_image(entry) == ("real.png", "s") + + +def test_extract_output_image_missing(): + with pytest.raises(ComfyUIError, match="no output images"): + extract_output_image({"outputs": {}}) diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..00410ac --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,32 @@ +from app.config import DEFAULT_MAPPING, Settings + + +def test_defaults(): + s = Settings(_env_file=None) + assert s.comfy_base_url == "http://192.168.0.3:8188" + assert s.execution_timeout == 300.0 + assert s.poll_interval == 1.5 + assert s.mapping() == DEFAULT_MAPPING + + +def test_env_override(monkeypatch): + monkeypatch.setenv("COMFY_BASE_URL", "http://example:9000") + monkeypatch.setenv("ADAPTER_API_KEY", "sk-abc") + monkeypatch.setenv("EXECUTION_TIMEOUT", "42") + monkeypatch.setenv("POLL_INTERVAL", "0.5") + s = Settings(_env_file=None) + assert s.comfy_base_url == "http://example:9000" + assert s.adapter_api_key == "sk-abc" + assert s.execution_timeout == 42.0 + assert s.poll_interval == 0.5 + + +def test_mapping_file_override(tmp_path): + f = tmp_path / "mapping.json" + f.write_text('{"create": {"prompt": {"node": "1", "field": "t"}}}', encoding="utf-8") + s = Settings(_env_file=None, mapping_file=f) + m = s.mapping() + assert m["create"]["prompt"] == {"node": "1", "field": "t"} + # other sections/keys keep defaults + assert m["edit"] == DEFAULT_MAPPING["edit"] + assert m["create"]["seed"] == DEFAULT_MAPPING["create"]["seed"] diff --git a/tests/test_edits.py b/tests/test_edits.py new file mode 100644 index 0000000..45213fe --- /dev/null +++ b/tests/test_edits.py @@ -0,0 +1,107 @@ +import base64 + +import httpx +import pytest +import respx + +from tests.conftest import COMFY, png_bytes, png_size + +EDITS = "http://adapter.test/v1/images/edits" +PROMPT_ID = "p-edit" +UPLOADED = "adapter_deadbeef_11223344.png" + +HISTORY = { + PROMPT_ID: { + "outputs": {"60": {"images": [{"filename": "ed_00001_.png", "subfolder": "", "type": "output"}]}}, + "status": {"status_str": "success", "completed": True}, + } +} + + +def _mock_pipeline(): + upload_route = respx.post(f"{COMFY}/upload/image").mock(return_value=httpx.Response(200, json={"name": UPLOADED, "type": "input"})) + prompt_route = respx.post(f"{COMFY}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": PROMPT_ID})) + respx.get(f"{COMFY}/history/{PROMPT_ID}").mock(return_value=httpx.Response(200, json=HISTORY)) + respx.get(f"{COMFY}/view").mock(return_value=httpx.Response(200, content=png_bytes(1024, 1024))) + return upload_route, prompt_route + + +@respx.mock +async def test_edits_multipart_success(client: httpx.AsyncClient): + upload_route, prompt_route = _mock_pipeline() + files = {"image": ("in.png", png_bytes(64, 64), "image/png")} + r = await client.post(EDITS, files=files, data={"prompt": "make it night"}) + assert r.status_code == 200 + body = r.json() + img = base64.b64decode(body["data"][0]["b64_json"]) + assert png_size(img) == (1024, 1024) + # uploaded file name landed in LoadImage node; prompt in positive TextEncode + upload_body = upload_route.calls.last.request.content + assert b"in.png" in upload_body or b"adapter_" in upload_body + sent = prompt_route.calls.last.request.content.decode() + assert UPLOADED in sent + assert "make it night" in sent + graph = prompt_route.calls.last.request.content # json bytes + assert b'"115:110"' in graph # negative node present with empty prompt + + +@respx.mock +async def test_edits_json_base64_success(client: httpx.AsyncClient): + _, prompt_route = _mock_pipeline() + b64 = base64.b64encode(png_bytes(64, 64)).decode() + r = await client.post(EDITS, json={"image": b64, "prompt": "add a hat", "model": "qwen-image-edit"}) + assert r.status_code == 200 + assert "add a hat" in prompt_route.calls.last.request.content.decode() + + +@respx.mock +async def test_edits_size_ignored(client: httpx.AsyncClient): + _, prompt_route = _mock_pipeline() + b64 = base64.b64encode(png_bytes(64, 64)).decode() + r = await client.post(EDITS, json={"image": b64, "prompt": "x", "size": "512x512"}) + assert r.status_code == 200 + sent = prompt_route.calls.last.request.content.decode() + assert "512" not in sent # size never reaches the workflow + + +async def test_edits_missing_image_400(client): + r = await client.post(EDITS, data={"prompt": "x"}) + assert r.status_code == 400 + + +async def test_edits_empty_image_400(client): + files = {"image": ("in.png", b"", "image/png")} + r = await client.post(EDITS, files=files, data={"prompt": "x"}) + assert r.status_code == 400 + + +async def test_edits_garbage_image_400(client): + files = {"image": ("in.png", b"not-an-image", "image/png")} + r = await client.post(EDITS, files=files, data={"prompt": "x"}) + assert r.status_code == 400 + + +async def test_edits_bad_base64_400(client): + r = await client.post(EDITS, json={"image": "!!!not base64!!!", "prompt": "x"}) + assert r.status_code == 400 + + +async def test_edits_empty_prompt_400(client): + b64 = base64.b64encode(png_bytes(8, 8)).decode() + r = await client.post(EDITS, json={"image": b64, "prompt": " "}) + assert r.status_code == 400 + + +async def test_edits_wrong_model_400(client): + b64 = base64.b64encode(png_bytes(8, 8)).decode() + r = await client.post(EDITS, json={"image": b64, "prompt": "x", "model": "z-image-turbo"}) + assert r.status_code == 400 + + +@respx.mock +async def test_edits_comfy_failure_502(client): + respx.post(f"{COMFY}/upload/image").mock(return_value=httpx.Response(200, json={"name": UPLOADED})) + respx.post(f"{COMFY}/prompt").mock(return_value=httpx.Response(400, json={"node_errors": {"x": {"errors": [{"message": "bad graph", "details": ""}]}}})) + b64 = base64.b64encode(png_bytes(8, 8)).decode() + r = await client.post(EDITS, json={"image": b64, "prompt": "x"}) + assert r.status_code == 502 diff --git a/tests/test_errors.py b/tests/test_errors.py new file mode 100644 index 0000000..d686e2d --- /dev/null +++ b/tests/test_errors.py @@ -0,0 +1,29 @@ +from app.errors import ComfyUIError, InvalidRequest, Unauthorized, error_body +from app.schemas import ImageGenerationRequest + +import pytest + + +def test_error_body_canonical_shape(): + body = error_body("boom", "server_error", "comfyui_timeout") + assert body == {"error": {"message": "boom", "type": "server_error", "param": None, "code": "comfyui_timeout"}} + + +def test_exception_status_mapping(): + assert InvalidRequest("x").status == 400 + assert Unauthorized().status == 401 + assert ComfyUIError("x").status == 502 + + +def test_generation_request_defaults(): + r = ImageGenerationRequest(prompt="p") + assert r.size == "1024x1024" + assert r.n == 1 + assert r.seed is None + assert r.model is None + + +@pytest.mark.parametrize("bad", ["", "abc", "1024", "1024x", "0x10", "10x0"]) +def test_size_validation(bad): + with pytest.raises(ValueError): + ImageGenerationRequest(prompt="p", size=bad) diff --git a/tests/test_generations.py b/tests/test_generations.py new file mode 100644 index 0000000..b5ad82a --- /dev/null +++ b/tests/test_generations.py @@ -0,0 +1,140 @@ +import base64 +import json + +import httpx +import pytest +import respx + +from tests.conftest import COMFY, png_bytes, png_size + +GEN = "http://adapter.test/v1/images/generations" +PROMPT_ID = "p-1" + + +def _history_json(width=64, height=64): + return { + PROMPT_ID: { + "outputs": {"60": {"images": [{"filename": "out_00001_.png", "subfolder": "", "type": "output"}]}}, + "status": {"status_str": "success", "completed": True}, + } + } + + +def _mock_pipeline(view_png: bytes | None = None, prompt_response: httpx.Response | None = None): + prompt_route = respx.post(f"{COMFY}/prompt").mock( + return_value=prompt_response or httpx.Response(200, json={"prompt_id": PROMPT_ID}) + ) + respx.get(f"{COMFY}/history/{PROMPT_ID}").mock(return_value=httpx.Response(200, json=_history_json())) + view_response = httpx.Response(200, content=view_png) if view_png else httpx.Response(500) + respx.get(f"{COMFY}/view").mock(return_value=view_response) + return prompt_route + + +@respx.mock +async def test_generation_success(client: httpx.AsyncClient): + png = png_bytes(1024, 1024) + prompt_route = _mock_pipeline(view_png=png) + r = await client.post(GEN, json={"prompt": "a red cube", "model": "z-image-turbo", "size": "1024x1024", "seed": 123}) + assert r.status_code == 200 + body = r.json() + assert set(body) == {"created", "data"} and set(body["data"][0]) == {"b64_json"} + img = base64.b64decode(body["data"][0]["b64_json"]) + assert png_size(img) == (1024, 1024) + # injected values reached the workflow + sent = json.loads(prompt_route.calls.last.request.content.decode())["prompt"] + assert sent["83:27"]["inputs"]["text"] == "a red cube" + assert sent["83:13"]["inputs"]["width"] == 1024 and sent["83:13"]["inputs"]["height"] == 1024 + assert sent["83:3"]["inputs"]["seed"] == 123 + + +@respx.mock +async def test_generation_size_custom(client: httpx.AsyncClient): + prompt_route = _mock_pipeline(view_png=png_bytes(832, 1216)) + r = await client.post(GEN, json={"prompt": "x", "size": "832x1216"}) + assert r.status_code == 200 + sent = json.loads(prompt_route.calls.last.request.content.decode())["prompt"] + assert (sent["83:13"]["inputs"]["width"], sent["83:13"]["inputs"]["height"]) == (832, 1216) + + +@respx.mock +async def test_generation_seed_random_per_request(client: httpx.AsyncClient): + prompt_route = _mock_pipeline(view_png=png_bytes()) + await client.post(GEN, json={"prompt": "x"}) + seed_a = json.loads(prompt_route.calls.last.request.content.decode())["prompt"]["83:3"]["inputs"]["seed"] + await client.post(GEN, json={"prompt": "x"}) + seed_b = json.loads(prompt_route.calls.last.request.content.decode())["prompt"]["83:3"]["inputs"]["seed"] + assert seed_a != seed_b + + +async def _assert_400(r): + assert r.status_code == 400 + err = r.json()["error"] + assert err["type"] == "invalid_request_error" and err["message"] + + +async def test_empty_prompt_400(client): + await _assert_400(await client.post(GEN, json={"prompt": " "})) + + +async def test_n_gt_1_400(client): + await _assert_400(await client.post(GEN, json={"prompt": "x", "n": 2})) + + +@pytest.mark.parametrize("bad", ["wide", "1024", "1024x", "0x512"]) +async def test_bad_size_400(client, bad): + await _assert_400(await client.post(GEN, json={"prompt": "x", "size": bad})) + + +async def test_unknown_model_400(client): + await _assert_400(await client.post(GEN, json={"prompt": "x", "model": "dall-e-3"})) + + +async def test_bad_json_400(client): + await _assert_400(await client.post(GEN, content=b"{not json", headers={"Content-Type": "application/json"})) + + +@respx.mock +async def test_comfy_validation_failure_502(client): + _mock_pipeline(prompt_response=httpx.Response(400, json={"node_errors": {"83:27": {"errors": [{"message": "bad", "details": "'83:43'"}]}}})) + r = await client.post(GEN, json={"prompt": "x"}) + assert r.status_code == 502 + assert r.json()["error"]["code"] == "comfyui_validation_failed" + + +@respx.mock +async def test_comfy_execution_failure_502(client): + respx.post(f"{COMFY}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": PROMPT_ID})) + respx.get(f"{COMFY}/history/{PROMPT_ID}").mock( + return_value=httpx.Response(200, json={PROMPT_ID: {"outputs": {}, "status": {"status_str": "error", "completed": False}}}) + ) + r = await client.post(GEN, json={"prompt": "x"}) + assert r.status_code == 502 + assert r.json()["error"]["code"] == "comfyui_execution_failed" + + +@respx.mock +async def test_comfy_timeout_502(client): + respx.post(f"{COMFY}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": PROMPT_ID})) + respx.get(f"{COMFY}/history/{PROMPT_ID}").mock(return_value=httpx.Response(200, json={})) + # client fixture uses execution_timeout=5; shorten via a second app instance instead: + import contextlib + + from app.main import create_app + from tests.conftest import settings as make_settings + + app = create_app(make_settings(execution_timeout=0.05, poll_interval=0.0)) + async with contextlib.AsyncExitStack() as stack: + await stack.enter_async_context(app.router.lifespan_context(app)) + transport = httpx.ASGITransport(app=app) + ac = await stack.enter_async_context(httpx.AsyncClient(transport=transport, base_url="http://adapter.test")) + ac.headers["Authorization"] = client.headers["Authorization"] + r = await ac.post(GEN, json={"prompt": "x"}) + assert r.status_code == 502 + assert r.json()["error"]["code"] == "comfyui_timeout" + + +@respx.mock +async def test_comfy_unreachable_502(client): + respx.post(f"{COMFY}/prompt").mock(side_effect=httpx.ConnectError("boom")) + r = await client.post(GEN, json={"prompt": "x"}) + assert r.status_code == 502 diff --git a/tests/test_injector.py b/tests/test_injector.py new file mode 100644 index 0000000..4d8a8cf --- /dev/null +++ b/tests/test_injector.py @@ -0,0 +1,73 @@ +import pytest + +from app.config import DEFAULT_MAPPING +from app.errors import InvalidRequest +from app.injector import build_workflow, new_seed, parse_size +from app.templates import load_all +from app.config import Settings + + +@pytest.fixture(scope="module") +def templates(): + return load_all(Settings(_env_file=None)) + + +def test_parse_size(): + assert parse_size("1024x1024") == (1024, 1024) + assert parse_size("832X1216") == (832, 1216) + + +@pytest.mark.parametrize("bad", ["1024", "ax1024", "1024x", "0x100", "1024x-2", "wide"]) +def test_parse_size_invalid(bad): + with pytest.raises(InvalidRequest): + parse_size(bad) + + +def test_seed_random_differs(): + assert new_seed(None) != new_seed(None) + + +def test_seed_fixed_and_range(): + assert new_seed(42) == 42 + with pytest.raises(InvalidRequest): + new_seed(-1) + + +def test_create_injection(templates): + create, _ = templates + mapping = DEFAULT_MAPPING["create"] + graph = build_workflow(create, mapping, prompt="a cat", seed=77, size="832x1216") + assert graph["83:27"]["inputs"]["text"] == "a cat" + assert graph["83:13"]["inputs"]["width"] == 832 + assert graph["83:13"]["inputs"]["height"] == 1216 + assert graph["83:3"]["inputs"]["seed"] == 77 + # template untouched + assert create["83:27"]["inputs"]["text"] == "\n" + assert create["83:13"]["inputs"]["width"] == 1280 + + +def test_edit_injection(templates): + _, edit = templates + mapping = DEFAULT_MAPPING["edit"] + graph = build_workflow(edit, mapping, prompt="make it night", seed=5, image_name="adapter_x.png") + assert graph["115:111"]["inputs"]["prompt"] == "make it night" + assert graph["115:110"]["inputs"]["prompt"] == "" + assert graph["78"]["inputs"]["image"] == "adapter_x.png" + assert graph["115:3"]["inputs"]["seed"] == 5 + + +def test_edit_seed_varies_per_call(templates): + _, edit = templates + mapping = DEFAULT_MAPPING["edit"] + seeds = set() + for _ in range(3): + g = build_workflow(edit, mapping, prompt="p", seed=new_seed(None)) + seeds.add(g["115:3"]["inputs"]["seed"]) + assert len(seeds) == 3 + + +def test_unknown_mapping_node_raises(templates): + create, _ = templates + bad_mapping = {"prompt": {"node": "999", "field": "text"}} + with pytest.raises(InvalidRequest, match="999"): + build_workflow(create, bad_mapping, prompt="p", seed=1) diff --git a/tests/test_live.py b/tests/test_live.py new file mode 100644 index 0000000..a8908df --- /dev/null +++ b/tests/test_live.py @@ -0,0 +1,70 @@ +"""Live tests against the real ComfyUI server (192.168.0.3) via a running adapter. + +Manual run: + 1. start adapter: uvicorn app.main:app --port 8000 + 2. set env: ADAPTER_LIVE=1 ADAPTER_URL=http://127.0.0.1:8000 ADAPTER_KEY= + 3. pytest -m live +""" +import base64 +import os +import struct + +import httpx +import pytest + +pytestmark = pytest.mark.live + +URL = os.environ.get("ADAPTER_URL", "http://127.0.0.1:8000") +KEY = os.environ.get("ADAPTER_KEY", "") + +pytestmark = [ + pytest.mark.live, + pytest.mark.skipif(not os.environ.get("ADAPTER_LIVE"), reason="set ADAPTER_LIVE=1 to run live tests"), +] + + +def _client() -> httpx.AsyncClient: + return httpx.AsyncClient(base_url=URL, headers={"Authorization": f"Bearer {KEY}"}, timeout=600) + + +def _png_size(data: bytes) -> tuple[int, int]: + return struct.unpack(">II", data[16:24]) + + +def _solid_png(w: int = 128, h: int = 96) -> bytes: + import zlib + + def chunk(t, d): + return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d)) + + ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0) + raw = b"".join(b"\x00" + bytes((200, 40, 40)) * w for _ in range(h)) + return b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", ihdr) + chunk(b"IDAT", zlib.compress(raw)) + chunk(b"IEND", b"") + + +async def test_generate_1024_with_seed(): + async with _client() as c: + r = await c.post("/v1/images/generations", json={"prompt": "a red cube on white background", "size": "1024x1024", "seed": 20261004}) + assert r.status_code == 200, r.text + img = base64.b64decode(r.json()["data"][0]["b64_json"]) + assert img.startswith(b"\x89PNG") + assert _png_size(img) == (1024, 1024) + + +async def test_same_seed_same_image(): + async with _client() as c: + payload = {"prompt": "a blue sphere", "size": "1024x1024", "seed": 777} + r1 = await c.post("/v1/images/generations", json=payload) + r2 = await c.post("/v1/images/generations", json=payload) + assert r1.json()["data"][0]["b64_json"] == r2.json()["data"][0]["b64_json"] + + +async def test_edit_output_is_about_1mp_with_input_aspect(): + async with _client() as c: + files = {"image": ("in.png", _solid_png(128, 96), "image/png")} + r = await c.post("/v1/images/edits", files=files, data={"prompt": "make the background deep blue"}) + assert r.status_code == 200, r.text + img = base64.b64decode(r.json()["data"][0]["b64_json"]) + w, h = _png_size(img) + assert 0.9 <= w * h / 1_000_000 <= 1.15 # ~1 MP + assert abs(w / h - 128 / 96) < 0.05 # aspect preserved diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000..14121f0 --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,11 @@ +import httpx + + +async def test_models_list(client: httpx.AsyncClient): + r = await client.get("/v1/models") + assert r.status_code == 200 + body = r.json() + assert body["object"] == "list" + ids = [m["id"] for m in body["data"]] + assert ids == ["z-image-turbo", "qwen-image-edit"] + assert all(m["object"] == "model" for m in body["data"]) diff --git a/tests/test_templates.py b/tests/test_templates.py new file mode 100644 index 0000000..465c5b4 --- /dev/null +++ b/tests/test_templates.py @@ -0,0 +1,57 @@ +import pytest + +from app.config import Settings +from app.templates import TemplateError, check_mapping, find_broken_links, load_all + + +def test_find_broken_links_detects_missing_node(): + graph = { + "1": {"class_type": "X", "inputs": {"m": ["9", 0]}}, + } + errors = find_broken_links(graph) + assert len(errors) == 1 + assert "missing node '9'" in errors[0] + assert "node 1" in errors[0] + + +def test_find_broken_links_ok_on_valid_graph(): + graph = { + "1": {"class_type": "X", "inputs": {"m": ["2", 0]}}, + "2": {"class_type": "Y", "inputs": {}}, + } + assert find_broken_links(graph) == [] + + +def test_check_mapping_missing_node(): + graph = {"1": {"class_type": "X", "inputs": {"text": ""}}} + mapping = {"create": {"prompt": {"node": "404", "field": "text"}}} + errors = check_mapping(graph, mapping["create"], "create") + assert any("not found" in e for e in errors) + + +def test_check_mapping_non_scalar_field(): + graph = {"1": {"class_type": "X", "inputs": {"text": ["2", 0]}}} + mapping = {"create": {"prompt": {"node": "1", "field": "text"}}} + errors = check_mapping(graph, mapping["create"], "create") + assert any("no scalar input" in e for e in errors) + + +def test_shipped_templates_load_cleanly(): + s = Settings(_env_file=None) + create, edit = load_all(s) + assert create["83:28"]["class_type"] == "UNETLoader" + assert edit["115:89"]["class_type"] == "LoraLoaderModelOnly" + + +def test_broken_template_blocks_startup(tmp_path): + bad = tmp_path / "bad.json" + bad.write_text('{"1": {"class_type": "X", "inputs": {"m": ["9", 0]}}}', encoding="utf-8") + good = tmp_path / "good.json" + good.write_text( + '{"1": {"class_type": "X", "inputs": {"m": ["2", 0]}}, "2": {"class_type": "Y", "inputs": {}}}', + encoding="utf-8", + ) + s = Settings(_env_file=None, create_template_path=bad, edit_template_path=good) + with pytest.raises(TemplateError) as exc: + load_all(s) + assert "missing node '9'" in str(exc.value) diff --git a/workflow_create.json b/workflow_create.json new file mode 100644 index 0000000..60615d0 --- /dev/null +++ b/workflow_create.json @@ -0,0 +1,78 @@ +{ + "60": { + "inputs": { + "filename_prefix": "z-image-turbo", + "images": ["83:8", 0] + }, + "class_type": "SaveImage", + "_meta": { "title": "Save Image" } + }, + "83:30": { + "inputs": { + "clip_name": "qwen_3_4b.safetensors", + "type": "lumina2", + "device": "default" + }, + "class_type": "CLIPLoader", + "_meta": { "title": "Load CLIP" } + }, + "83:29": { + "inputs": { "vae_name": "ae.safetensors" }, + "class_type": "VAELoader", + "_meta": { "title": "Load VAE" } + }, + "83:13": { + "inputs": { + "width": 1280, + "height": 1280, + "batch_size": 1 + }, + "class_type": "EmptySD3LatentImage", + "_meta": { "title": "EmptySD3LatentImage" } + }, + "83:33": { + "inputs": { "conditioning": ["83:27", 0] }, + "class_type": "ConditioningZeroOut", + "_meta": { "title": "ConditioningZeroOut" } + }, + "83:8": { + "inputs": { + "samples": ["83:3", 0], + "vae": ["83:29", 0] + }, + "class_type": "VAEDecode", + "_meta": { "title": "VAE Decode" } + }, + "83:3": { + "inputs": { + "seed": 578351132403975, + "steps": 9, + "cfg": 1, + "sampler_name": "res_multistep", + "scheduler": "simple", + "denoise": 1, + "model": ["83:28", 0], + "positive": ["83:27", 0], + "negative": ["83:33", 0], + "latent_image": ["83:13", 0] + }, + "class_type": "KSampler", + "_meta": { "title": "KSampler" } + }, + "83:27": { + "inputs": { + "text": "\n", + "clip": ["83:30", 0] + }, + "class_type": "CLIPTextEncode", + "_meta": { "title": "CLIP Text Encode (Prompt)" } + }, + "83:28": { + "inputs": { + "unet_name": "z_image_turbo_bf16.safetensors", + "weight_dtype": "default" + }, + "class_type": "UNETLoader", + "_meta": { "title": "Load Diffusion Model" } + } +} diff --git a/workflow_edit.json b/workflow_edit.json new file mode 100644 index 0000000..248b411 --- /dev/null +++ b/workflow_edit.json @@ -0,0 +1,138 @@ +{ + "60": { + "inputs": { + "filename_prefix": "ComfyUI", + "images": ["115:8", 0] + }, + "class_type": "SaveImage", + "_meta": { "title": "Save Image" } + }, + "78": { + "inputs": { "image": "" }, + "class_type": "LoadImage", + "_meta": { "title": "Load Image" } + }, + "115:75": { + "inputs": { + "strength": 1, + "model": ["115:66", 0] + }, + "class_type": "CFGNorm", + "_meta": { "title": "CFGNorm" } + }, + "115:39": { + "inputs": { "vae_name": "qwen_image_vae.safetensors" }, + "class_type": "VAELoader", + "_meta": { "title": "Load VAE" } + }, + "115:38": { + "inputs": { + "clip_name": "qwen_2.5_vl_7b_fp8_scaled.safetensors", + "type": "qwen_image", + "device": "default" + }, + "class_type": "CLIPLoader", + "_meta": { "title": "Load CLIP" } + }, + "115:37": { + "inputs": { + "unet_name": "qwen_image_edit_2509_fp8_e4m3fn.safetensors", + "weight_dtype": "default" + }, + "class_type": "UNETLoader", + "_meta": { "title": "Load Diffusion Model" } + }, + "115:110": { + "inputs": { + "prompt": "", + "clip": ["115:38", 0], + "vae": ["115:39", 0], + "image1": ["115:93", 0] + }, + "class_type": "TextEncodeQwenImageEditPlus", + "_meta": { "title": "TextEncodeQwenImageEditPlus" } + }, + "115:66": { + "inputs": { + "shift": 3, + "model": ["115:89", 0] + }, + "class_type": "ModelSamplingAuraFlow", + "_meta": { "title": "ModelSamplingAuraFlow" } + }, + "115:111": { + "inputs": { + "prompt": "", + "clip": ["115:38", 0], + "vae": ["115:39", 0], + "image1": ["115:93", 0] + }, + "class_type": "TextEncodeQwenImageEditPlus", + "_meta": { "title": "TextEncodeQwenImageEditPlus" } + }, + "115:112": { + "inputs": { + "width": 2560, + "height": 1440, + "batch_size": 1 + }, + "class_type": "EmptySD3LatentImage", + "_meta": { "title": "EmptySD3LatentImage" } + }, + "115:88": { + "inputs": { + "pixels": ["115:93", 0], + "vae": ["115:39", 0] + }, + "class_type": "VAEEncode", + "_meta": { "title": "VAE Encode" } + }, + "115:93": { + "inputs": { + "upscale_method": "lanczos", + "megapixels": 1, + "resolution_steps": 1, + "image": ["78", 0] + }, + "class_type": "ImageScaleToTotalPixels", + "_meta": { "title": "ImageScaleToTotalPixels" } + }, + "115:3": { + "inputs": { + "seed": 1118877715456453, + "steps": 4, + "cfg": 1, + "sampler_name": "euler", + "scheduler": "simple", + "denoise": 1, + "model": ["115:75", 0], + "positive": ["115:111", 0], + "negative": ["115:110", 0], + "latent_image": ["115:88", 0] + }, + "class_type": "KSampler", + "_meta": { "title": "KSampler" } + }, + "115:8": { + "inputs": { + "samples": ["115:3", 0], + "vae": ["115:39", 0] + }, + "class_type": "VAEDecode", + "_meta": { "title": "VAE Decode" } + }, + "115:116": { + "inputs": { "images": ["115:8", 0] }, + "class_type": "PreviewImage", + "_meta": { "title": "PreviewImage" } + }, + "115:89": { + "inputs": { + "lora_name": "Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16.safetensors", + "strength_model": 1, + "model": ["115:37", 0] + }, + "class_type": "LoraLoaderModelOnly", + "_meta": { "title": "LoraLoaderModelOnly" } + } +}