From 0c5c0b423b5b2a8393f5830cd6817845b5a0937f Mon Sep 17 00:00:00 2001 From: Shaman Date: Mon, 5 Oct 2026 02:19:46 +0300 Subject: [PATCH] =?UTF-8?q?=D0=94=D0=BE=D1=80=D0=B0=D0=B1=D0=BE=D1=82?= =?UTF-8?q?=D0=BA=D0=B0=20=D1=83=D0=BF=D1=80=D0=B0=D0=B2=D0=BB=D0=B5=D0=BD?= =?UTF-8?q?=D0=B8=D1=8F=20=D0=BC=D0=BE=D0=B4=D0=B5=D0=BB=D1=8F=D0=BC=D0=B8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env.example | 5 +- Dockerfile | 3 +- README.md | 166 ++++++++++++------ app/config.py | 64 +++---- app/errors.py | 7 + app/injector.py | 109 ++++++++---- app/main.py | 61 ++++--- app/templates.py | 112 ++++++++---- models.json | 10 ++ tests/test_config.py | 53 ++++-- tests/test_generations.py | 5 + tests/test_injector.py | 65 ++++--- tests/test_registry_integration.py | 101 +++++++++++ tests/test_templates.py | 104 ++++++++--- .../qwen_image_edit.json | 8 +- .../z_image_turbo.json | 8 +- 16 files changed, 629 insertions(+), 252 deletions(-) create mode 100644 models.json create mode 100644 tests/test_registry_integration.py rename workflow_edit.json => workflows/qwen_image_edit.json (96%) rename workflow_create.json => workflows/z_image_turbo.json (94%) diff --git a/.env.example b/.env.example index 2217c2e..409715e 100644 --- a/.env.example +++ b/.env.example @@ -6,5 +6,6 @@ ADAPTER_API_KEY=sk-1234 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 +# Model registry: JSON mapping an API model id to its workflow file (required). +# Defaults to models.json next to the app; override the path here if you move it. +MODELS_FILE=models.json diff --git a/Dockerfile b/Dockerfile index 2eee778..3e8a371 100644 --- a/Dockerfile +++ b/Dockerfile @@ -5,7 +5,8 @@ ENV PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1 COPY pyproject.toml ./ COPY app ./app -COPY workflow_create.json workflow_edit.json ./ +COPY models.json ./ +COPY workflows ./workflows RUN pip install --no-cache-dir . EXPOSE 8000 diff --git a/README.md b/README.md index 474bbcf..2b6d2b5 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,13 @@ **Модели** -| id модели в API | Назначение | Диффузионная модель | Доп. | +В API различаются два понятия «модель». **API-модель** — внешний алиас из реестра +`models.json`, который лишь выбирает, какой workflow выполнить (см. «Реестр моделей»). +**Модель ComfyUI** (диффузионный чекпойнт, LoRA, CLIP, VAE) неизменяема и зашита внутрь +самого workflow-графа — она не подставляется из запроса. Каждая API-модель жёстко привязана +к своему workflow-файлу с уже подобранной цепочкой моделей. + +| id модели в API | Назначение | Диффузионная модель (внутри workflow) | Доп. | |---|---|---|---| | `z-image-turbo` | генерация по текстовому промпту | `z_image_turbo_bf16.safetensors` | CLIP `qwen_3_4b` (lumina2), VAE `ae` | | `qwen-image-edit` | редактирование входного изображения | `qwen_image_edit_2509_fp8_e4m3fn.safetensors` | LoRA `Qwen-Image-Edit-2509-Lightning-4steps-V1.0-bf16` (4 шага), CLIP `qwen_2.5_vl_7b_fp8_scaled` (qwen_image), VAE `qwen_image_vae` | @@ -37,7 +43,7 @@ - **`GET /v1/models`** — список доступных «моделей» (для автодополнения в клиентах). - **Воспроизводимость**: опциональный `seed`; один и тот же `seed` + промпт дают идентичный результат. - **Толерантный парсинг OpenAI**: неприменимые поля (`quality`, `style`, `user`) не ломают запрос. -- **Workflow как данные**: шаблоны JSON и карта маппинга параметров на узлы вынесены из кода. +- **Workflow как данные**: реестр моделей `models.json` и workflow-файлы с плейсхолдерами вынесены из кода. - **Защита от деградации графа**: битые ссылки в шаблоне обнаруживаются локально и на сервере, адаптер не стартует с нерабочим workflow. - **Статический API-ключ** на всех эндпоинтах, кроме `GET /health`. - **Ошибки в формате OpenAI** (`error.message/type/param/code`) — клиенты корректно их разбирают. @@ -87,10 +93,8 @@ curl -s http://localhost:8000/health | `ADAPTER_API_KEY` | пусто (**обязательно**) | Ожидаемый Bearer-ключ клиентов. Пустое значение = все запросы получают 401 (fail-closed) | | `EXECUTION_TIMEOUT` | `300` | Секунд ожидания результата от ComfyUI; превышение -> 502 `comfyui_timeout` | | `POLL_INTERVAL` | `1.5` | Интервал опроса `/history`, сек | -| `CREATE_TEMPLATE_PATH` | `<репо>/workflow_create.json` | Путь к шаблону генерации | -| `EDIT_TEMPLATE_PATH` | `<репо>/workflow_edit.json` | Путь к шаблону редактирования | -| `MAPPING_FILE` | не задан | JSON с переопределением карты маппинга (см. ниже) | -| `VALIDATE_TEMPLATES_ON_SERVER` | `true` | Прогонять графы через `POST /prompt?front=false` при старте | +| `MODELS_FILE` | `<репо>/models.json` | Реестр моделей: JSON `id -> {workflow, owned_by}` (обязателен, см. «Реестр моделей») | +| `VALIDATE_TEMPLATES_ON_SERVER` | `true` | Прогонять собранные графы через `POST /prompt?front=false` при старте | Логирование стартовых проверок видно с `uvicorn --log-level info`: @@ -121,7 +125,7 @@ Authorization: Bearer | Параметр | Тип | По умолчанию | Обязателен | Описание | |---|---|---|---|---| | `prompt` | string | — | **да** | Текстовое описание. Пустая строка -> 400 | -| `model` | string \| null | `null` | нет | Если задан — только `z-image-turbo`, иначе 400 | +| `model` | string \| null | `null` | нет | Id модели из реестра; пусто -> первая совместимая с генерацией (`z-image-turbo`). Неизвестная/несовместимая -> 400 | | `n` | int \| null | `1` | нет | Поддерживается только `1`, иначе 400 | | `size` | string | `1024x1024` | нет | Строго `WxH`, оба значения > 0. Попала в latent-ноду, то есть реально задаёт разрешение | | `seed` | int \| null | `null` | нет | `0..2^63-1`. Без него — криптографически случайный (`secrets.randbelow`) | @@ -146,7 +150,7 @@ Authorization: Bearer |---|---|---|---| | `image` | файл | **да** | PNG или JPEG (проверка по сигнатуре байтов) | | `prompt` | string | **да** | Инструкция по редактированию | -| `model` | string | нет | Только `qwen-image-edit` | +| `model` | string | нет | Id из реестра, совместимый с edits; пусто -> первая такая (`qwen-image-edit`). Иначе 400 | | `seed` | int | нет | В форме передаётся строкой, парсится в int | **application/json:** @@ -155,7 +159,7 @@ Authorization: Bearer |---|---|---|---| | `image` | string (base64) | **да** | Допускается data-URL-префикс `data:image/png;base64,...` | | `prompt` | string | **да** | | -| `model` | string | нет | Только `qwen-image-edit` | +| `model` | string | нет | Id из реестра, совместимый с edits; пусто -> первая такая (`qwen-image-edit`). Иначе 400 | | `seed` | int | нет | | | `size` | string | нет | **Принимается, но не влияет на результат** — см. «Ограничения» | @@ -283,7 +287,7 @@ request [1] auth + валидация тела --> 400 / 401 | v -[2] deep-copy шаблона workflow, подстановка prompt/size/seed/image по карте маппинга +[2] deep-copy workflow выбранной модели, замещение плейсхолдеров {prompt}/{seed}/{width}/{height}/{image} значениями запроса | v [3] (только edits) POST /upload/image -> имя файла в input-каталоге ComfyUI @@ -306,49 +310,98 @@ WebSocket `/ws` намеренно не используется: ComfyUI гар соединение и устойчив к разрывам. Транспорт изолирован в `app/comfy_client.py` — переход на websocket не затронул бы остальной код. -### Карта маппинга +### Реестр моделей (models.json) -Соответствие параметра запроса узлу workflow. Правится **без изменения кода** через -`MAPPING_FILE` (JSON с секциями `create` и `edit`; заданные ключи перезаписывают -дефолт, остальные сохраняются). - -| Параметр | `workflow_create.json` | `workflow_edit.json` | -|---|---|---| -| `prompt` | `83:27` -> `inputs.text` (CLIPTextEncode) | `115:111` -> `inputs.prompt` (TextEncodeQwenImageEditPlus) | -| negative | — (в графе `ConditioningZeroOut`, константа) | `115:110` -> `inputs.prompt` (всегда `""`) | -| `size` | `83:13` -> `inputs.width` / `inputs.height` (EmptySD3LatentImage) | — (см. «Ограничения») | -| `seed` | `83:3` -> `inputs.seed` (KSampler) | `115:3` -> `inputs.seed` (KSampler) | -| `image` | — | `78` -> `inputs.image` (LoadImage) | - -Пример переопределения: +Список API-моделей задаётся **декларативно**, в файле `models.json` (путь — env +`MODELS_FILE`, по умолчанию `<репо>/models.json`). Каждая запись связывает id модели с её +workflow-файлом: ```json -{ "create": { "prompt": { "node": "99:1", "field": "text" } } } +{ + "z-image-turbo": { "workflow": "workflows/z_image_turbo.json", "owned_by": "comfy-adapter" }, + "qwen-image-edit": { "workflow": "workflows/qwen_image_edit.json", "owned_by": "comfy-adapter" } +} ``` -При старте каждый id из карты проверяется на существование в шаблоне и на то, что -поле скалярное, — рассинхрон карты и шаблона падает на старт, а не на первый запрос. +- **Добавить модель = добавить запись + файл workflow. Код править не нужно.** +- `workflow` — путь к JSON (относительный разрешается относительно папки `models.json`). +- `owned_by` необязателен (по умолчанию `comfy-adapter`); используется в `GET /v1/models`. +- **Порядок записей значим**: он задаёт приоритет выбора при `model` не задан. +- Файл обязателен: отсутствующий или битый `models.json` -> явная ошибка конфигурации на + старте (fail-fast), сервис не поднимется молча. -### Проверка шаблонов +**Выбор модели.** Поле `model` выбирает запись реестра. Если `model` не задан или `null` — +берётся **первая совместимая с эндпоинтом** модель. Неизвестный id или несовместимый с +эндпоинтом -> `400`. Совместимость вычисляется из набора токенов (см. ниже), отдельного +поля «тип модели» в реестре нет. -Двухуровневая, оба уровня на старте: +**Совместимость модели с эндпоинтом** — вложенность множеств токенов, наборы принадлежат +эндпоинту: -1. **Локально** (`app/templates.py`): обход всех `inputs` и поиск ссылок - `[node_id, slot]` на несуществующие узлы. Найденное -> исключение со списком - `node X: input 'Y' links to missing node 'Z'`, адаптер не стартует. -2. **На сервере** (`VALIDATE_TEMPLATES_ON_SERVER=true`): `POST /prompt?front=false` — - ComfyUI валидирует граф, **не ставя задание в очередь** и не тратя GPU. +``` + модель допустима для эндпоинта <=> required(эндп) ⊆ токены(граф) ⊆ supplied(эндп) + + эндпоинт supplied (что может подставить) required (обязательно в графе) + generations prompt, negative, seed, width, height prompt + edits prompt, negative, seed, image prompt, image +``` + +Отсюда: генерация не примет модель с `{image}` (не может подставить), а редактирование не +примет модель без `{image}` (картинка не должна остаться незамещённой). + +### Плейсхолдеры в workflow + +Точка внедрения значений находится **внутри workflow-графа** — в значениях полей узлов +пишутся токены вида `{имя}`. Никакой внешней карты узел↔поле нет; соответствие задаёт сам +граф. Фиксированный словарь токенов принадлежит адаптеру: + +| Токен | Тип | Кто подставляет | +|---|---|---| +| `{prompt}` | строка | оба эндпоинта | +| `{negative}` | строка | всегда `""` (запрос негатив не принимает) | +| `{seed}` | число | оба эндпоинта | +| `{width}` / `{height}` | число | только генерация (из `size`) | +| `{image}` | строка | только edits (имя загруженного файла) | + +Пример узла KSampler с плейсхолдерами: + +```json +"83:3": { "class_type": "KSampler", "inputs": { "seed": "{seed}", "...": 0, "model": ["83:28", 0] } } +``` + +Правила подстановки: + +- значение узла **равно ровно** `{seed}` -> подставляется **число** (типы сохраняются; + иначе `front=false` не примет строку в int-виджете); токен внутри текста -> строковая + вставка без искажения окружения; +- нераспознанный токен (вне словаря) -> workflow отклоняется на старте; +- незамещённый токен (эндпоинт не поставил значение) -> запрос отклоняется, ComfyUI не + получает задание. + +### Проверка workflow + +Двухуровневая, оба уровня на старте, для **каждой** модели реестра: + +1. **Локально** (`app/templates.py`): обход `inputs` в поиске ссылок `[node_id, slot]` на + несуществующие узлы и проверка токенов (нет ли вне словаря). Найденное -> исключение со + списком `node X: input 'Y' links to missing node 'Z'`, адаптер не стартует. +2. **На сервере** (`VALIDATE_TEMPLATES_ON_SERVER=true`): граф каждой модели собирается + типизированными значениями-заглушками (dry-inject: `seed=1`, `width=height=1024`, + `prompt="x"`, `negative=""`, `image="dummy.png"`) и отправляется на + `POST /prompt?front=false` — ComfyUI валидирует граф, **не ставя задание в очередь** и + не тратя GPU. Заодно прогоняется сам инжектор. Ошибка логируется как диагностика старта + и не блокирует запуск (согласно прежнему поведению). --- -## Шаблоны workflow и их починка +## Workflow-файлы и их починка -Оба шаблона были выгружены из OpenWebUI (там прямая интеграция с ComfyUI). В +Оба workflow-файла были выгружены из OpenWebUI (там прямая интеграция с ComfyUI). В **исходном виде они не проходили валидацию ComfyUI** — в цепочках моделей отсутствовали узлы, на которые шли ссылки: ``` -workflow_create.json workflow_edit.json +workflows/z_image_turbo.json workflows/qwen_image_edit.json 83:28 UNETLoader 115:37 UNETLoader | на него никто не ссылается | на него никто не ссылается @@ -359,12 +412,13 @@ workflow_create.json workflow_edit.json Ошибка сервера: `prompt_outputs_failed_validation`, `KeyError: '83:43'` / `'115:117'`. -Исправления, применённые в репозиторных файлах (минимальные, по одной ссылке на шаблон): +Исправления, применённые в репозиторных файлах (минимальные, по одной ссылке на workflow; +в те же файлы добавлены плейсхолдеры `{prompt}`/`{seed}`/`{width}`/`{height}`/`{image}`/`{negative}`): -| Шаблон | Было | Стало | +| Workflow | Было | Стало | |---|---|---| -| `workflow_create.json` | `83:3.inputs.model = ["83:43", 0]`, `83:27.inputs.clip = ["83:43", 1]` | `["83:28", 0]` (UNETLoader), `["83:30", 0]` (CLIPLoader) | -| `workflow_edit.json` | `115:66.inputs.model = ["115:117", 0]` | `["115:89", 0]` — LoRA-нода `115:89` подключена в цепочку (`115:37 -> 115:89 -> 115:66 -> 115:75 -> KSampler`) | +| `workflows/z_image_turbo.json` | `83:3.inputs.model = ["83:43", 0]`, `83:27.inputs.clip = ["83:43", 1]` | `["83:28", 0]` (UNETLoader), `["83:30", 0]` (CLIPLoader) | +| `workflows/qwen_image_edit.json` | `115:66.inputs.model = ["115:117", 0]` | `["115:89", 0]` — LoRA-нода `115:89` подключена в цепочку (`115:37 -> 115:89 -> 115:66 -> 115:75 -> KSampler`) | Подключение LoRA во втором случае содержательно верно: в настройках OpenWebUI в поле «Модель» редактирования указан именно Lightning-LoRA на 4 шага, а `steps: 4` в @@ -376,7 +430,7 @@ KSampler без него не имеет смысла. При смене workflow в OpenWebUI следует перезалить JSON сюда и сверить эти две цепочки; локальная проверка на старте поймает проблему автоматически. -### Свойства шаблонов, которые стоит знать +### Свойства workflow, которые стоит знать - **Размер выхода в edit ≠ `size`.** В графе за разрешение отвечает `ImageScaleToTotalPixels` (`megapixels: 1`), а latent из `EmptySD3LatentImage` @@ -396,10 +450,10 @@ KSampler без него не имеет смысла. comfy-adapter/ ├── app/ │ ├── main.py FastAPI: эндпоинты, цикл execute, разбор multipart/JSON -│ ├── config.py настройки из env + дефолтная карта маппинга +│ ├── config.py настройки из env + загрузка реестра моделей из models.json │ ├── comfy_client.py async-клиент ComfyUI: /prompt, /upload/image, /history, /view -│ ├── injector.py deep-copy шаблона + подстановка значений, парсинг WxH, seed -│ ├── templates.py загрузка шаблонов, проверка битых ссылок и соответствия карте +│ ├── injector.py deep-copy workflow + замещение плейсхолдеров, парсинг WxH, seed +│ ├── templates.py загрузка реестра, проверка ссылок и токенов, выбор модели │ ├── schemas.py pydantic-модели OpenAI-контракта │ ├── errors.py иерархия исключений + обработчики в формат OpenAI │ └── auth.py проверка статического Bearer-ключа @@ -408,15 +462,16 @@ comfy-adapter/ │ ├── test_generations.py генерация: успех, инъекция значений, seed, все классы 400/502 │ ├── test_edits.py edits: multipart и base64, игнорирование size, ошибки входа │ ├── test_comfy_client.py каждый метод клиента на respx-моках (успех/HTTP-ошибка/таймаут) -│ ├── test_injector.py маппинг, deep-copy без мутации шаблона, seed, парсинг size -│ ├── test_templates.py битые ссылки, сверка с картой, валидность репозиторных шаблонов -│ ├── test_config.py дефолты, переопределение из env, слияние MAPPING_FILE +│ ├── test_injector.py плейсхолдеры, deep-copy без мутации, seed, парсинг size +│ ├── test_templates.py битые ссылки, токены, совместимость и выбор модели, ошибки реестра +│ ├── test_config.py дефолты, переопределение из env, загрузка и fail-fast реестра │ ├── test_auth.py ключ: отсутствует / неверный / верный; health без ключа │ ├── test_errors.py канонический вид error-body, соответствие статусов │ ├── test_models.py формат /v1/models +│ ├── test_registry_integration.py новая модель только правкой models.json (без кода) │ └── test_live.py живые прогоны против реального ComfyUI (маркер live) -├── workflow_create.json исправленный шаблон генерации -├── workflow_edit.json исправленный шаблон редактирования +├── models.json реестр API-моделей: id -> {workflow, owned_by} +├── workflows/ workflow-файлы с плейсхолдерами (z_image_turbo.json, qwen_image_edit.json) ├── Dockerfile python:3.12-slim + uvicorn ├── pyproject.toml └── .env.example @@ -450,11 +505,14 @@ seed (повтор даёт идентичный base64), редактирова ## Ограничения -- **`n > 1` не поддерживается** -> 400. Нужен — правьте `batch_size` в шаблоне или +- **`n > 1` не поддерживается** -> 400. Нужен — правьте `batch_size` в workflow или запускайте несколько запросов. -- **`size` в `/v1/images/edits` не влияет на выход** (см. «Свойства шаблонов»). -- **`negative_prompt` из запроса не принимается.** Узел в шаблоне есть, но поля в - схеме запроса нет — адаптер всегда пишет пустую строку. +- **`size` в `/v1/images/edits` не влияет на выход** (см. «Свойства workflow»): edit-граф не + содержит `{width}`/`{height}`, эндпоинт edits их не поставляет. +- **`negative_prompt` из запроса не принимается.** Поля в схеме запроса нет — токен + `{negative}`, если он есть в workflow, всегда резолвится в пустую строку `""`. +- **`model` не задан -> берётся первая совместимая с эндпоинтом модель** реестра (порядок + записей в `models.json` задаёт приоритет). Явный `default`-флаг пока не реализован. - **Размер в генерации не квантуется.** Любое корректное `WxH` уходит в latent напрямую; нестандартные значения могут дать артефакты или ошибку уже на стороне ComfyUI (прилетит как 502, а не 400). diff --git a/app/config.py b/app/config.py index dc5d37e..681b43b 100644 --- a/app/config.py +++ b/app/config.py @@ -1,30 +1,20 @@ -"""Environment settings and the parameter -> workflow node mapping.""" +"""Environment settings and the model registry loaded from a JSON file.""" import json from pathlib import Path +from pydantic import BaseModel, ValidationError from pydantic_settings import BaseSettings, SettingsConfigDict +from .errors import ConfigError + 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 ModelEntry(BaseModel): + """A registry record: an API model id maps to one workflow file.""" + + workflow: Path + owned_by: str = "comfy-adapter" class Settings(BaseSettings): @@ -34,21 +24,31 @@ class Settings(BaseSettings): 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. + # JSON file mapping an API model id to its workflow. Required; no defaults in code. + models_file: Path = REPO_ROOT / "models.json" + # Validate collected 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 load_registry(self) -> dict[str, ModelEntry]: + """Load the model registry, preserving declaration order (first = preferred). + + Fail-fast: a missing or malformed file raises ConfigError and blocks startup. + """ + try: + raw = json.loads(self.models_file.read_text(encoding="utf-8")) + except FileNotFoundError as e: + raise ConfigError(f"model registry file not found: {self.models_file}") from e + except json.JSONDecodeError as e: + raise ConfigError(f"model registry is not valid JSON ({self.models_file}): {e}") from e + if not isinstance(raw, dict) or not raw: + raise ConfigError(f"model registry must be a non-empty object of model entries: {self.models_file}") + registry: dict[str, ModelEntry] = {} + for model_id, spec in raw.items(): + try: + registry[model_id] = ModelEntry.model_validate(spec) + except ValidationError as e: + raise ConfigError(f"model '{model_id}' has an invalid entry: {e}") from e + return registry def get_settings() -> Settings: diff --git a/app/errors.py b/app/errors.py index e8394c8..0b6379d 100644 --- a/app/errors.py +++ b/app/errors.py @@ -31,6 +31,13 @@ class ComfyUIError(AdapterError): super().__init__(502, message, "server_error", code) +class ConfigError(AdapterError): + """Invalid model registry configuration; blocks startup.""" + + def __init__(self, message: str): + super().__init__(500, message, "server_error", "config_error") + + def error_body(message: str, err_type: str, code: str | None = None) -> dict: return {"error": {"message": message, "type": err_type, "param": None, "code": code}} diff --git a/app/injector.py b/app/injector.py index c543aae..3988529 100644 --- a/app/injector.py +++ b/app/injector.py @@ -1,5 +1,12 @@ -"""Build a concrete workflow graph from a template + request parameters.""" +"""Build a concrete workflow graph from a template + request values. + +Values are injected through placeholder tokens written directly inside the +workflow graph (e.g. "{prompt}", "{seed}"), so the graph itself declares which +node input receives which request field. The set of recognised tokens is fixed +and owned by the adapter, not by the model. +""" import copy +import re import secrets from typing import Any @@ -7,6 +14,21 @@ from .errors import InvalidRequest SEED_MAX = 2**63 - 1 +# Fixed token vocabulary owned by the adapter. A model's workflow may only use +# these tokens; anything else is rejected. Types drive typed vs string injection. +TOKEN_TYPES: dict[str, str] = { + "prompt": "str", + "negative": "str", + "seed": "int", + "width": "int", + "height": "int", + "image": "str", +} +KNOWN_TOKENS = frozenset(TOKEN_TYPES) + +_TOKEN = re.compile(r"\{(\w+)\}") +_WHOLE = re.compile(r"^\{(\w+)\}$") + def parse_size(size: str) -> tuple[int, int]: """Parse OpenAI `WxH` size strings (e.g. '1024x1024').""" @@ -28,40 +50,55 @@ def new_seed(seed: int | None) -> int: 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.""" +def extract_tokens(graph: dict[str, Any]) -> tuple[set[str], list[str]]: + """Token names referenced by the graph. + + Returns (known_tokens, unknown_token_errors). A token is any `{name}` inside + a string node input. Unknown names are reported so startup can reject the + workflow before it ever reaches a request. + """ + declared: set[str] = set() + unknown: set[str] = set() + for node_id, node in graph.items(): + for field, value in node.get("inputs", {}).items(): + if isinstance(value, str): + for name in _TOKEN.findall(value): + if name in KNOWN_TOKENS: + declared.add(name) + else: + unknown.add(f"node {node_id}: input '{field}' uses unknown token '{{{name}}}'") + return declared, sorted(unknown) + + +def build_workflow(template: dict[str, Any], values: dict[str, Any]) -> dict[str, Any]: + """Deep-copy the template and substitute placeholder tokens with values. + + Injection is type-preserving: a value that is exactly one token (e.g. + "{seed}") is replaced by the typed object (int for seed/width/height, str + otherwise); a token embedded in surrounding text is string-interpolated. + A referenced token that is unrecognized or missing from ``values`` rejects + the build, so a malformed graph never reaches ComfyUI. + """ 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 + for node in graph.values(): + inputs = node.get("inputs", {}) + for field, value in inputs.items(): + if isinstance(value, str): + inputs[field] = _substitute(value, values) return graph + + +def _substitute(value: str, values: dict[str, Any]) -> Any: + for name in _TOKEN.findall(value): + if name not in KNOWN_TOKENS: + raise InvalidRequest(f"workflow uses unrecognized token '{{{name}}}'") + if name not in values: + raise InvalidRequest(f"no value supplied for token '{{{name}}}'") + whole = _WHOLE.fullmatch(value) + if whole: + return _coerce(whole.group(1), values[whole.group(1)]) + return _TOKEN.sub(lambda m: str(values[m.group(1)]), value) + + +def _coerce(name: str, value: Any) -> Any: + return int(value) if TOKEN_TYPES[name] == "int" else str(value) diff --git a/app/main.py b/app/main.py index 8af1be7..05acbcf 100644 --- a/app/main.py +++ b/app/main.py @@ -8,9 +8,9 @@ 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 .config import Settings, get_settings from .errors import InvalidRequest, install_error_handlers -from .injector import build_workflow, new_seed +from .injector import build_workflow, new_seed, parse_size from .schemas import ( ImageEditJsonRequest, ImageGenerationRequest, @@ -19,32 +19,42 @@ from .schemas import ( ModelCard, ModelsResponse, ) -from .templates import load_all +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): - create_graph, edit_graph = load_all(settings) + registry = load_registry(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) + 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.templates = {"create": create_graph, "edit": edit_graph} - logger.info("comfy-adapter ready (server=%s)", settings.comfy_base_url) + app.state.registry = registry + logger.info("comfy-adapter ready (server=%s, models=%s)", settings.comfy_base_url, list(registry)) try: yield finally: @@ -75,7 +85,10 @@ def create_app(settings: Settings | None = None) -> FastAPI: @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]) + 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: @@ -86,14 +99,11 @@ def create_app(settings: Settings | None = None) -> FastAPI: 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}'") + model = select_model(request.app.state.registry, body.model, "generations") + width, height = parse_size(body.size) graph = build_workflow( - app.state.templates["create"], - settings.mapping()["create"], - prompt=body.prompt, - seed=new_seed(body.seed), - size=body.size, + model.graph, + {"prompt": body.prompt, "seed": new_seed(body.seed), "negative": "", "width": width, "height": height}, ) return await _run(graph) @@ -112,18 +122,13 @@ def create_app(settings: Settings | None = None) -> FastAPI: 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) + model = select_model(request.app.state.registry, model, "edits") 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="", + model.graph, + {"prompt": prompt, "seed": new_seed(seed), "negative": "", "image": name}, ) return await _run(graph) diff --git a/app/templates.py b/app/templates.py index 332bd0e..6d2fe2d 100644 --- a/app/templates.py +++ b/app/templates.py @@ -1,9 +1,17 @@ -"""Template loading with startup validation of graph links and mappings.""" +"""Model registry loading: read each workflow, validate links/tokens, and +decide which model an endpoint may run. + +Startup validation of a workflow is local here (broken links + token check); +the collected graphs are then dry-injected and sent to the server elsewhere +(see main.lifespan). +""" import json +from dataclasses import dataclass, field from pathlib import Path from .config import Settings -from .errors import AdapterError +from .errors import AdapterError, ConfigError, InvalidRequest +from .injector import extract_tokens class TemplateError(AdapterError): @@ -13,6 +21,29 @@ class TemplateError(AdapterError): super().__init__(500, message, "server_error", "template_error") +@dataclass(frozen=True) +class Model: + id: str + graph: dict + declared: frozenset[str] + owned_by: str = "comfy-adapter" + workflow: Path = field(default=Path(".")) + + +# Token sets belong to the endpoint, not the model. A model is valid for an +# endpoint iff required(endpoint) <= declared(graph) <= supplied(endpoint). +ENDPOINT_VOCAB: dict[str, dict[str, frozenset[str]]] = { + "generations": { + "supplied": frozenset({"prompt", "negative", "seed", "width", "height"}), + "required": frozenset({"prompt"}), + }, + "edits": { + "supplied": frozenset({"prompt", "negative", "seed", "image"}), + "required": frozenset({"prompt", "image"}), + }, +} + + def find_broken_links(graph: dict) -> list[str]: """References to node ids that do not exist in the graph.""" errors = [] @@ -24,39 +55,58 @@ def find_broken_links(graph: dict) -> list[str]: 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 is_compatible(declared: frozenset[str], endpoint: str) -> bool: + vocab = ENDPOINT_VOCAB[endpoint] + return vocab["required"] <= declared and declared <= vocab["supplied"] -def _mapped_fields(spec: dict) -> list[str]: - if "width_field" in spec: - return [spec["width_field"], spec["height_field"]] - return [spec["field"]] +def select_model(models: dict[str, Model], model: str | None, endpoint: str) -> Model: + """Pick the model for a request by id, or the first compatible one. + + An explicit id that is unknown or incompatible with the endpoint is a 400, + as is a missing id when no registered model is compatible. + """ + if model is not None: + found = models.get(model) + if found is None: + raise InvalidRequest(f"unknown model '{model}'") + if not is_compatible(found.declared, endpoint): + raise InvalidRequest(f"model '{model}' is not compatible with this endpoint") + return found + for candidate in models.values(): # declaration order = priority + if is_compatible(candidate.declared, endpoint): + return candidate + raise InvalidRequest(f"no model in the registry is compatible with this endpoint") -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) +def _read_graph(path: Path, model_id: str) -> dict: + try: + with open(path, encoding="utf-8") as f: + return json.load(f) + except FileNotFoundError as e: + raise TemplateError(f"model '{model_id}': workflow file not found: {path}") from e + except json.JSONDecodeError as e: + raise TemplateError(f"model '{model_id}': workflow is not valid JSON ({path}): {e}") from e + + +def load_model(model_id: str, workflow_path: Path, owned_by: str) -> Model: + graph = _read_graph(workflow_path, model_id) + errors = find_broken_links(graph) + declared, unknown = extract_tokens(graph) + errors += unknown if errors: - raise TemplateError(f"workflow template '{graph_name}' ({path}) is invalid:\n " + "\n ".join(errors)) - return graph + raise TemplateError( + f"workflow for model '{model_id}' ({workflow_path}) is invalid:\n " + "\n ".join(errors) + ) + return Model(id=model_id, graph=graph, declared=frozenset(declared), owned_by=owned_by, workflow=workflow_path) -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 +def load_registry(settings: Settings) -> dict[str, Model]: + """Load every model in the registry, preserving order for default selection.""" + base_dir = settings.models_file.parent + registry = settings.load_registry() + models: dict[str, Model] = {} + for model_id, entry in registry.items(): + path = entry.workflow if entry.workflow.is_absolute() else base_dir / entry.workflow + models[model_id] = load_model(model_id, path, entry.owned_by) + return models diff --git a/models.json b/models.json new file mode 100644 index 0000000..c5a3b3f --- /dev/null +++ b/models.json @@ -0,0 +1,10 @@ +{ + "z-image-turbo": { + "workflow": "workflows/z_image_turbo.json", + "owned_by": "comfy-adapter" + }, + "qwen-image-edit": { + "workflow": "workflows/qwen_image_edit.json", + "owned_by": "comfy-adapter" + } +} diff --git a/tests/test_config.py b/tests/test_config.py index 00410ac..822ef60 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,4 +1,7 @@ -from app.config import DEFAULT_MAPPING, Settings +import pytest + +from app.config import Settings +from app.errors import ConfigError def test_defaults(): @@ -6,7 +9,7 @@ def test_defaults(): 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 + assert s.models_file.name == "models.json" def test_env_override(monkeypatch): @@ -21,12 +24,40 @@ def test_env_override(monkeypatch): 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"] +def test_registry_loads_shipped_in_order(): + s = Settings(_env_file=None) + registry = s.load_registry() + # order preserved (first = preferred for default selection) + assert list(registry) == ["z-image-turbo", "qwen-image-edit"] + assert registry["z-image-turbo"].workflow.name == "z_image_turbo.json" + assert registry["z-image-turbo"].owned_by == "comfy-adapter" + + +def test_registry_missing_file_fails_fast(tmp_path): + s = Settings(_env_file=None, models_file=tmp_path / "nope.json") + with pytest.raises(ConfigError, match="not found"): + s.load_registry() + + +def test_registry_malformed_json_fails_fast(tmp_path): + f = tmp_path / "models.json" + f.write_text("{not valid", encoding="utf-8") + s = Settings(_env_file=None, models_file=f) + with pytest.raises(ConfigError, match="not valid JSON"): + s.load_registry() + + +def test_registry_empty_fails_fast(tmp_path): + f = tmp_path / "models.json" + f.write_text("{}", encoding="utf-8") + s = Settings(_env_file=None, models_file=f) + with pytest.raises(ConfigError, match="non-empty"): + s.load_registry() + + +def test_registry_bad_entry_fails_fast(tmp_path): + f = tmp_path / "models.json" + f.write_text('{"m": {"not_workflow": "x.json"}}', encoding="utf-8") + s = Settings(_env_file=None, models_file=f) + with pytest.raises(ConfigError, match="invalid entry"): + s.load_registry() diff --git a/tests/test_generations.py b/tests/test_generations.py index b5ad82a..910422e 100644 --- a/tests/test_generations.py +++ b/tests/test_generations.py @@ -89,6 +89,11 @@ async def test_unknown_model_400(client): await _assert_400(await client.post(GEN, json={"prompt": "x", "model": "dall-e-3"})) +async def test_incompatible_model_400(client): + # qwen-image-edit declares {image}, which generations cannot supply + await _assert_400(await client.post(GEN, json={"prompt": "x", "model": "qwen-image-edit"})) + + async def test_bad_json_400(client): await _assert_400(await client.post(GEN, content=b"{not json", headers={"Content-Type": "application/json"})) diff --git a/tests/test_injector.py b/tests/test_injector.py index 4d8a8cf..38676ea 100644 --- a/tests/test_injector.py +++ b/tests/test_injector.py @@ -1,15 +1,14 @@ import pytest -from app.config import DEFAULT_MAPPING +from app.config import Settings 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 +from app.templates import load_registry @pytest.fixture(scope="module") -def templates(): - return load_all(Settings(_env_file=None)) +def registry(): + return load_registry(Settings(_env_file=None)) def test_parse_size(): @@ -33,41 +32,57 @@ def test_seed_fixed_and_range(): 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") +def test_create_injection_and_template_untouched(registry): + create = registry["z-image-turbo"].graph + graph = build_workflow(create, {"prompt": "a cat", "seed": 77, "negative": "", "width": 832, "height": 1216}) 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 + # template still holds placeholders (deep copy, no mutation) + assert create["83:27"]["inputs"]["text"] == "{prompt}" + assert create["83:13"]["inputs"]["width"] == "{width}" + assert create["83:3"]["inputs"]["seed"] == "{seed}" -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") +def test_injection_is_typed(registry): + graph = build_workflow(registry["z-image-turbo"].graph, {"prompt": "x", "seed": 7, "negative": "", "width": 64, "height": 128}) + assert isinstance(graph["83:3"]["inputs"]["seed"], int) # numeric token -> int, not "{seed}" + assert isinstance(graph["83:13"]["inputs"]["width"], int) + assert isinstance(graph["83:27"]["inputs"]["text"], str) + + +def test_embedded_token_is_string_interpolated(): + graph = {"1": {"class_type": "X", "inputs": {"text": "prefix {prompt} suffix"}}} + out = build_workflow(graph, {"prompt": "hi"}) + assert out["1"]["inputs"]["text"] == "prefix hi suffix" + + +def test_edit_injection(registry): + edit = registry["qwen-image-edit"].graph + graph = build_workflow(edit, {"prompt": "make it night", "seed": 5, "negative": "", "image": "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"] +def test_edit_seed_varies_per_call(registry): + edit = registry["qwen-image-edit"].graph seeds = set() for _ in range(3): - g = build_workflow(edit, mapping, prompt="p", seed=new_seed(None)) + g = build_workflow(edit, {"prompt": "p", "seed": new_seed(None), "negative": "", "image": "a.png"}) 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) +def test_unrecognized_token_rejected(): + graph = {"1": {"class_type": "X", "inputs": {"text": "{foo}"}}} + with pytest.raises(InvalidRequest, match="unrecognized token"): + build_workflow(graph, {"prompt": "p"}) + + +def test_unfilled_token_rejected(): + graph = {"1": {"class_type": "X", "inputs": {"seed": "{seed}"}}} + with pytest.raises(InvalidRequest, match="no value supplied"): + build_workflow(graph, {"prompt": "p"}) diff --git a/tests/test_registry_integration.py b/tests/test_registry_integration.py new file mode 100644 index 0000000..9686a7b --- /dev/null +++ b/tests/test_registry_integration.py @@ -0,0 +1,101 @@ +"""A third model added purely by editing models.json (no code change): it shows up +in /v1/models, is dry-inject validated at startup, and is selectable by id.""" +import contextlib +import json + +import httpx +import respx + +from app.config import REPO_ROOT, Settings +from app.main import create_app +from tests.conftest import COMFY, FAKE_KEY + +PROMPT_ID = "p-int" + + +def _extra_generation_workflow() -> dict: + # declared = {prompt, negative, seed, width, height} -> generations-compatible, no {image} + return { + "1": { + "class_type": "X", + "inputs": { + "text": "{prompt}", + "neg": "{negative}", + "sd": "{seed}", + "w": "{width}", + "h": "{height}", + }, + } + } + + +def _models_file(tmp_path) -> Settings: + (tmp_path / "extra.json").write_text(json.dumps(_extra_generation_workflow()), encoding="utf-8") + registry = { + "z-image-turbo": {"workflow": str(REPO_ROOT / "workflows" / "z_image_turbo.json")}, + "extra-gen": {"workflow": "extra.json"}, # relative to models.json dir + "qwen-image-edit": {"workflow": str(REPO_ROOT / "workflows" / "qwen_image_edit.json")}, + } + models = tmp_path / "models.json" + models.write_text(json.dumps(registry), encoding="utf-8") + return Settings( + _env_file=None, + comfy_base_url=COMFY, + adapter_api_key=FAKE_KEY, + models_file=models, + validate_templates_on_server=True, + poll_interval=0.0, + execution_timeout=5.0, + ) + + +async def _with_lifespan(app): + stack = contextlib.AsyncExitStack() + 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"] = f"Bearer {FAKE_KEY}" + return stack, ac + + +@respx.mock +async def test_third_model_listed_and_validated_without_code_change(tmp_path): + validate_route = respx.post(f"{COMFY}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": PROMPT_ID})) + app = create_app(_models_file(tmp_path)) + stack, ac = await _with_lifespan(app) + try: + ids = [m["id"] for m in (await ac.get("/v1/models")).json()["data"]] + assert ids == ["z-image-turbo", "extra-gen", "qwen-image-edit"] + # every model dry-injected and validated once at startup + assert validate_route.call_count == 3 + finally: + await stack.aclose() + + +@respx.mock +async def test_extra_model_selectable_and_injects(tmp_path): + 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={PROMPT_ID: {"outputs": {"1": {"images": [{"filename": "o.png", "subfolder": "", "type": "output"}]}}, + "status": {"status_str": "success", "completed": True}}}, + ) + ) + respx.get(f"{COMFY}/view").mock(return_value=httpx.Response(200, content=b"\x89PNG\r\n\x1a\n")) + + settings = _models_file(tmp_path).model_copy(update={"validate_templates_on_server": False}) + app = create_app(settings) + stack, ac = await _with_lifespan(app) + try: + r = await ac.post( + "http://adapter.test/v1/images/generations", + json={"model": "extra-gen", "prompt": "hello", "size": "512x768", "seed": 9}, + ) + assert r.status_code == 200 + sent = json.loads(prompt_route.calls.last.request.content.decode())["prompt"] + assert sent["1"]["inputs"]["text"] == "hello" + assert sent["1"]["inputs"]["sd"] == 9 + assert sent["1"]["inputs"]["w"] == 512 and sent["1"]["inputs"]["h"] == 768 + finally: + await stack.aclose() diff --git a/tests/test_templates.py b/tests/test_templates.py index 465c5b4..a4c075e 100644 --- a/tests/test_templates.py +++ b/tests/test_templates.py @@ -1,7 +1,15 @@ import pytest from app.config import Settings -from app.templates import TemplateError, check_mapping, find_broken_links, load_all +from app.errors import InvalidRequest +from app.injector import extract_tokens +from app.templates import ( + TemplateError, + find_broken_links, + is_compatible, + load_registry, + select_model, +) def test_find_broken_links_detects_missing_node(): @@ -22,36 +30,84 @@ def test_find_broken_links_ok_on_valid_graph(): 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_extract_tokens_reports_known_and_unknown(): + graph = { + "1": {"class_type": "X", "inputs": {"text": "{prompt}", "seed": "{seed}"}}, + "2": {"class_type": "Y", "inputs": {"note": "static {notatoken} text"}}, + } + declared, unknown = extract_tokens(graph) + assert declared == {"prompt", "seed"} + assert len(unknown) == 1 and "notatoken" in unknown[0] -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_registry_loads_with_expected_tokens(): + registry = load_registry(Settings(_env_file=None)) + assert set(registry) == {"z-image-turbo", "qwen-image-edit"} + assert registry["z-image-turbo"].declared == frozenset({"prompt", "seed", "width", "height"}) + assert registry["qwen-image-edit"].declared == frozenset({"prompt", "negative", "seed", "image"}) -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_generations_graph_incompatible_with_edits(): + assert not is_compatible(frozenset({"prompt", "seed", "width", "height"}), "edits") -def test_broken_template_blocks_startup(tmp_path): +def test_edits_graph_incompatible_with_generations(): + assert not is_compatible(frozenset({"prompt", "negative", "seed", "image"}), "generations") + + +def test_select_model_defaults(registry): + assert select_model(registry, None, "generations").id == "z-image-turbo" + assert select_model(registry, None, "edits").id == "qwen-image-edit" + + +def test_select_model_explicit(registry): + assert select_model(registry, "qwen-image-edit", "edits").id == "qwen-image-edit" + + +def test_select_model_unknown_400(registry): + with pytest.raises(InvalidRequest, match="unknown model"): + select_model(registry, "dall-e-3", "generations") + + +def test_select_model_incompatible_400(registry): + with pytest.raises(InvalidRequest, match="not compatible"): + select_model(registry, "qwen-image-edit", "generations") + with pytest.raises(InvalidRequest, match="not compatible"): + select_model(registry, "z-image-turbo", "edits") + + +def test_select_model_no_suitable_400(registry): + gen_only = {"z-image-turbo": registry["z-image-turbo"]} + with pytest.raises(InvalidRequest, match="no model"): + select_model(gen_only, None, "edits") + + +def test_broken_workflow_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) + models = tmp_path / "models.json" + models.write_text('{"m1": {"workflow": "bad.json"}}', encoding="utf-8") with pytest.raises(TemplateError) as exc: - load_all(s) + load_registry(Settings(_env_file=None, models_file=models)) assert "missing node '9'" in str(exc.value) + + +def test_unknown_token_blocks_startup(tmp_path): + wf = tmp_path / "wf.json" + wf.write_text('{"1": {"class_type": "X", "inputs": {"text": "{bogus}"}}}', encoding="utf-8") + models = tmp_path / "models.json" + models.write_text('{"m1": {"workflow": "wf.json"}}', encoding="utf-8") + with pytest.raises(TemplateError, match="unknown token"): + load_registry(Settings(_env_file=None, models_file=models)) + + +def test_missing_workflow_file_blocks_startup(tmp_path): + models = tmp_path / "models.json" + models.write_text('{"m1": {"workflow": "absent.json"}}', encoding="utf-8") + with pytest.raises(TemplateError, match="not found"): + load_registry(Settings(_env_file=None, models_file=models)) + + +@pytest.fixture(scope="module") +def registry(): + return load_registry(Settings(_env_file=None)) diff --git a/workflow_edit.json b/workflows/qwen_image_edit.json similarity index 96% rename from workflow_edit.json rename to workflows/qwen_image_edit.json index 248b411..e9e3526 100644 --- a/workflow_edit.json +++ b/workflows/qwen_image_edit.json @@ -8,7 +8,7 @@ "_meta": { "title": "Save Image" } }, "78": { - "inputs": { "image": "" }, + "inputs": { "image": "{image}" }, "class_type": "LoadImage", "_meta": { "title": "Load Image" } }, @@ -44,7 +44,7 @@ }, "115:110": { "inputs": { - "prompt": "", + "prompt": "{negative}", "clip": ["115:38", 0], "vae": ["115:39", 0], "image1": ["115:93", 0] @@ -62,7 +62,7 @@ }, "115:111": { "inputs": { - "prompt": "", + "prompt": "{prompt}", "clip": ["115:38", 0], "vae": ["115:39", 0], "image1": ["115:93", 0] @@ -99,7 +99,7 @@ }, "115:3": { "inputs": { - "seed": 1118877715456453, + "seed": "{seed}", "steps": 4, "cfg": 1, "sampler_name": "euler", diff --git a/workflow_create.json b/workflows/z_image_turbo.json similarity index 94% rename from workflow_create.json rename to workflows/z_image_turbo.json index 60615d0..a9003c0 100644 --- a/workflow_create.json +++ b/workflows/z_image_turbo.json @@ -23,8 +23,8 @@ }, "83:13": { "inputs": { - "width": 1280, - "height": 1280, + "width": "{width}", + "height": "{height}", "batch_size": 1 }, "class_type": "EmptySD3LatentImage", @@ -45,7 +45,7 @@ }, "83:3": { "inputs": { - "seed": 578351132403975, + "seed": "{seed}", "steps": 9, "cfg": 1, "sampler_name": "res_multistep", @@ -61,7 +61,7 @@ }, "83:27": { "inputs": { - "text": "\n", + "text": "{prompt}", "clip": ["83:30", 0] }, "class_type": "CLIPTextEncode",