commit 6d3d2fca4652b5d045803cedadd7a1e8e0ee86d3 Author: Alexander Garzon Date: Mon Jun 8 09:39:21 2026 -0400 feat: initial public release of ComfyUI-Imagent BYOK OpenAI gpt-image generation and editing nodes for ComfyUI — calls go directly to the OpenAI API with your own key (no proxy, no credit markup). - Two nodes: "Imagent: OpenAI Image" (text-to-image) and "Imagent: OpenAI Image Edit" (edit / inpaint / multi-reference compose). - Models gpt-image-2 / gpt-image-1.5 / gpt-image-1 with a per-model dynamic UI (comfy_api DynamicCombo): size, background, custom WxH, and input_fidelity are revealed only for the models that support them. - Edit node accepts an auto-growing set of reference images (up to 16) plus an optional inpaint mask; generate supports batches up to 8. - API key via OPENAI_API_KEY or a local config.json. Requires ComfyUI >= 0.23.0. diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..70cff25 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,63 @@ +name: Publish to ComfyUI Registry + +on: + push: + tags: ["v*"] + workflow_dispatch: + inputs: + tag: + description: "Existing tag whose changelog to (re)sync, e.g. v0.1.0 (skips re-publish)" + required: true + +permissions: + contents: read + +jobs: + publish: + runs-on: ubuntu-latest + env: + COMFY_API_KEY: ${{ secrets.COMFY_API_KEY }} + PUBLISHER_ID: agarzon + NODE_ID: comfyui-imagent + steps: + - uses: actions/checkout@v6 + with: + ref: ${{ github.event.inputs.tag || github.ref_name }} + + - uses: actions/setup-python@v6 + with: + python-version: "3.12" + + - name: Publish to ComfyUI Registry + if: github.event_name == 'push' && env.COMFY_API_KEY != '' + run: | + pip install comfy-cli + comfy node publish --token "$COMFY_API_KEY" + + - name: Sync version changelog to the Registry + if: env.COMFY_API_KEY != '' + env: + TAG_INPUT: ${{ github.event.inputs.tag }} + run: | + TAG="${TAG_INPUT:-$GITHUB_REF_NAME}" + v="${TAG#v}" + awk -v v="$v" ' + $0 ~ "^## v?" v " " { f=1; next } + f && /^## / { exit } + f { print } + ' CHANGELOG.md > changelog.md + if [ ! -s changelog.md ]; then + echo "No CHANGELOG section found for $v" >&2 + exit 1 + fi + body=$(jq -Rs '{changelog: ., deprecated: false}' < changelog.md) + code=$(curl -sS -o resp.json -w '%{http_code}' -X PUT \ + "https://api.comfy.org/publishers/$PUBLISHER_ID/nodes/$NODE_ID/versions/$v" \ + -H "Authorization: Bearer $COMFY_API_KEY" \ + -H "Content-Type: application/json" \ + -d "$body") + echo "HTTP $code"; cat resp.json; echo + case "$code" in + 2*) echo "Changelog synced for $v" ;; + *) echo "Failed to sync changelog (HTTP $code)" >&2; exit 1 ;; + esac diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..2d75abf --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,45 @@ +name: Release + +on: + push: + tags: ["v*"] + workflow_dispatch: + inputs: + tag: + description: "Existing tag to (re)release, e.g. v0.1.0" + required: true + +permissions: + contents: write + +jobs: + release: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + + - name: Resolve tag and extract its CHANGELOG section + id: notes + env: + TAG_INPUT: ${{ github.event.inputs.tag }} + run: | + TAG="${TAG_INPUT:-$GITHUB_REF_NAME}" + v="${TAG#v}" + echo "tag=$TAG" >> "$GITHUB_OUTPUT" + awk -v v="$v" ' + $0 ~ "^## v?" v " " { f=1; next } + f && /^## / { exit } + f { print } + ' CHANGELOG.md > release-notes.md + if [ ! -s release-notes.md ]; then + echo "No CHANGELOG section found for $v" >&2 + exit 1 + fi + cat release-notes.md + + - name: Create GitHub Release + uses: softprops/action-gh-release@v3 + with: + tag_name: ${{ steps.notes.outputs.tag }} + name: ${{ steps.notes.outputs.tag }} + body_path: release-notes.md diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..170ca54 --- /dev/null +++ b/.gitignore @@ -0,0 +1,26 @@ +# Python +__pycache__/ +*.py[cod] +.venv/ +.pytest_cache/ +*.egg-info/ +build/ +dist/ + +# Editor / OS +.vscode/ +.idea/ +.DS_Store + +# Project-local (secrets) +config.json + +# Claude Code (local only) +.claude/ +CLAUDE.md + +# Design docs (kept out of the public repo) +docs/superpowers/ + +# Dev: sample files bind-mounted into the container +docker/sample-files/ diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..6d8e0d5 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,25 @@ +# Changelog + +## v0.1.0 — 2026-06-08 + +First public release. + +### Added + +- **Two nodes** — *Imagent: OpenAI Image* (text-to-image via `images.generate`) and + *Imagent: OpenAI Image Edit* (edit / inpaint / multi-reference compose via `images.edit`). +- **Models** — `gpt-image-2`, `gpt-image-1.5`, `gpt-image-1` (DALL·E excluded — shut down 2026-05-12). +- **Per-model dynamic UI** — the `model` widget swaps the size/background options per model and + reveals only what each model supports: hi-res presets and custom `WxH` sizes for gpt-image-2, + `transparent` background and `input_fidelity` for gpt-image-1.x. `custom_width` / `custom_height` + appear only when `size = custom`. +- **Batch + multi-reference** — up to 8 images per call; the edit node takes an auto-growing set + of up to 16 reference images plus an optional inpaint mask (white = region to edit). +- Quality, background, output format (`png` / `jpeg` / `webp`), output compression, and content + moderation controls. +- **BYOK** — calls the OpenAI API directly with your own key (env `OPENAI_API_KEY` or a local + `config.json`); no proxy, no per-image credit markup. + +### Requirements + +- ComfyUI ≥ 0.23.0 (the nodes use the `comfy_api` IO schema with `DynamicCombo`). diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..73be743 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Alexander Garzon + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..fcf2b87 --- /dev/null +++ b/README.md @@ -0,0 +1,141 @@ +

+ ComfyUI-Imagent — OpenAI image generation and editing for ComfyUI +

+ +

ComfyUI-Imagent

+ +

+ BYOK OpenAI image generation and editing nodes for ComfyUI.
+ Bring your own OpenAI API key — calls go directly to OpenAI, no proxy, no extra credits. +

+ +

+ License: MIT + Python 3.10+ + Sponsored by Wallrus +

+ +--- + +## ✨ Features + +- 🖼️ **Two nodes** — **Imagent: OpenAI Image** (text-to-image) and **Imagent: OpenAI Image Edit** (edit, inpaint, multi-reference compose). +- 🤖 **Current gpt-image models** — `gpt-image-2`, `gpt-image-1.5`, `gpt-image-1`. +- ✍️ **Text-to-image** with full control over size, quality, background, and output format. +- 🎨 **Edit + inpaint (mask)** — supply a mask to repaint a specific region; white pixels mark the area to edit. +- 🔗 **Multi-reference compositing** — feed up to 16 reference images to the edit node (auto-growing input). +- 🎛️ **Per-model UI** — the `model` widget is dynamic: switching models swaps the size list and shows only the options that model supports (no invalid combos). +- 🏞️ **Up to 8 images per call** — batch generation in a single node execution. +- 🛡️ **Optional content moderation** — `auto` (default) or `low`. +- 🔑 **BYOK — your own OpenAI API key, your own costs.** Unlike ComfyUI's built-in OpenAI nodes (which route through ComfyUI's paid proxy), Imagent calls the OpenAI API directly. No proxy, no per-image credit markup. +- 🚫 DALL·E intentionally excluded — all DALL·E models were shut down by OpenAI in 2026. + +## 📦 Installation + +> **Requires ComfyUI ≥ 0.23.0** — the nodes use the `comfy_api` IO schema (`DynamicCombo`) for the per-model UI. + +### Option A — ComfyUI Manager (recommended) +Open **ComfyUI Manager → Custom Nodes Manager**, search for **`Imagent`**, click **Install**, and restart ComfyUI. + +### Option B — git clone +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/agarzon/ComfyUI-Imagent.git +pip install -r ComfyUI-Imagent/requirements.txt +``` +Restart ComfyUI. + +## 🔑 API key + +Imagent reads your key in this order: + +1. **Environment variable** (preferred): set `OPENAI_API_KEY` before launching ComfyUI. +2. **`config.json`**: copy `config.example.json` → `config.json` (gitignored) in the extension directory and add your key. + +> **Note:** OpenAI may require **API Organization Verification** before gpt-image models can be called. If your key is valid but you receive an authorization error, complete verification in the OpenAI dashboard — the error message logged to the ComfyUI console will tell you. + +## 🧩 Nodes + +### 🤖 Imagent: OpenAI Image + +Text-to-image generation via `images.generate`. + +| Parameter | Type | Values / Notes | +|---|---|---| +| `prompt` | STRING | Text description of the image to generate | +| `model` | **DynamicCombo** | `gpt-image-2` (default), `gpt-image-1.5`, `gpt-image-1`. Switching the model swaps the options below. | +| ↳ `size` | COMBO | **gpt-image-2:** `auto`, the three base sizes, five high-res presets, and `custom`. **gpt-image-1.x:** `auto` + the three base sizes only. | +| ↳↳ `custom_width` / `custom_height` | INT | **Appear only when `size = custom`.** 1024–3840, step 16; both multiples of 16; aspect ≤ 3:1; total pixels 655,360–8,294,400. | +| ↳ `background` | COMBO | **gpt-image-2:** `auto`, `opaque`. **gpt-image-1.x:** `auto`, `opaque`, `transparent` (needs `png`/`webp`). | +| `quality` | COMBO | `auto`, `low`, `medium`, `high` | +| `output_format` | COMBO | `png`, `jpeg`, `webp` | +| `output_compression` | INT | 0–100. Applied only for `jpeg` and `webp`. | +| `moderation` | COMBO | `auto` (default), `low`. Sent only when not `auto`. | +| `n` | INT | 1–8 images per call | + +Rows marked ↳ live inside the dynamic `model` widget (appear only for models that support them); ↳↳ rows are nested one level deeper and appear only when their parent option is selected. + +> No `seed` widget: OpenAI's image API has no seed, so it can't make outputs reproducible. ComfyUI's normal input-based caching applies — re-queuing an unchanged graph returns the cached image; change the prompt (or any input) to regenerate. + +**Output:** `image` (IMAGE tensor, batch of n). Errors are logged to the ComfyUI console. + +--- + +### 🤖 Imagent: OpenAI Image Edit + +Image editing, inpainting, and multi-reference compositing via `images.edit`. + +| Parameter | Type | Values / Notes | +|---|---|---| +| `prompt` | STRING | Description of the desired edit | +| `model` | **DynamicCombo** | Same three models; switching swaps the options below. | +| ↳ `size` | COMBO | Same as the generate node (hi-res + `custom` on gpt-image-2; base sizes on gpt-image-1.x). | +| ↳↳ `custom_width` / `custom_height` | INT | **Appear only when `size = custom`** (gpt-image-2). Same rules as the generate node. | +| ↳ `background` | COMBO | gpt-image-2: `auto`/`opaque`; gpt-image-1.x: adds `transparent`. | +| ↳ `input_fidelity` | COMBO | **gpt-image-1.x only.** `high` (default) or `low` — how closely to follow the reference image. | +| `images` | IMAGE (auto-grow) | Reference image(s) to edit — grows up to **16** slots; at least one required. | +| `quality` | COMBO | `auto`, `low`, `medium`, `high` | +| `output_format` | COMBO | `png`, `jpeg`, `webp` | +| `moderation` | COMBO | `auto` (default), `low` | +| `n` | INT | 1–8 images per call | +| `mask` *(optional)* | MASK | **White = region to edit.** Inpainting requires exactly one reference image. | + +Rows marked ↳ live inside the dynamic `model` widget; ↳↳ rows appear only when their parent option is selected. + +**Output:** `image` (IMAGE tensor, batch of n). Errors are logged to the ComfyUI console. + +**Quick-start recipes:** +- **Edit:** connect one image to `images`, write a `prompt`, leave `mask` disconnected. +- **Inpaint:** connect one image + `mask` (white = area to repaint), write a `prompt`. +- **Multi-reference:** connect several images (the input grows as you wire them up), write a `prompt`. Mask not supported for multi-image calls. + +## 🛠️ Development + +```bash +python3 -m venv .venv +.venv/bin/pip install -e ".[test]" +.venv/bin/pytest -q +``` + +### Smoke testing with Docker +Set your API key first — either export `OPENAI_API_KEY` in your shell, or copy `config.example.json` → `config.json` and add your key. Then: + +```bash +docker compose -f docker/docker-compose.yml up -d --build +``` + +ComfyUI will be available at **http://localhost:8188** with Imagent pre-loaded. For Python changes, restart the container. The repo is bind-mounted into the container's `custom_nodes/` directory. + +## 🐾 The story behind the name + +**Imagent** = **image** + **agent**. A small model-driven agent that images things into existence — and edits them once they're there. + +It comes from **Wallrus**, whose own name blends a *social **wall*** with a ***walrus***. Two friendly ideas, one tool: straightforward, reliable image generation without the middleman. + +## 💙 Sponsored by Wallrus + +ComfyUI-Imagent is proudly sponsored by **[Wallrus](https://wallrus.tech)**. If Imagent makes your ComfyUI workflow nicer, go say hi. 👋 + +## 📄 License + +[MIT](LICENSE) © 2026 Alexander Garzon diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..914b7f7 --- /dev/null +++ b/__init__.py @@ -0,0 +1,33 @@ +"""ComfyUI-Imagent: OpenAI gpt-image generation and editing nodes. + +The nodes use the comfy_api IO schema (DynamicCombo for per-model widgets), so we +register them via the V3 `comfy_entrypoint`/ComfyExtension path. Requires +ComfyUI >= 0.23.0 (for IO.DynamicCombo). +""" +import os +import sys + +# ComfyUI loads this __init__.py via importlib.util.spec_from_file_location, +# which does NOT add the custom-node directory to sys.path. Add it so that +# `import imagent.*` resolves whether installed editable or bind-mounted. +_here = os.path.dirname(os.path.abspath(__file__)) +if _here not in sys.path: + sys.path.insert(0, _here) + +# The nodes use the comfy_api IO schema, which only exists inside ComfyUI. Guard +# the import so the package root still imports cleanly without ComfyUI (e.g. the +# dev venv / test collection); ComfyUI itself always provides comfy_api. +try: + from comfy_api.latest import ComfyExtension, IO # noqa: E402 +except ModuleNotFoundError: + ComfyExtension = None +else: + from imagent.node_edit import OpenAIImageEdit # noqa: E402 + from imagent.node_generate import OpenAIImageGenerate # noqa: E402 + + class ImagentExtension(ComfyExtension): + async def get_node_list(self) -> list[type[IO.ComfyNode]]: + return [OpenAIImageGenerate, OpenAIImageEdit] + + async def comfy_entrypoint() -> ImagentExtension: + return ImagentExtension() diff --git a/assets/banner.png b/assets/banner.png new file mode 100644 index 0000000..7cbc58c Binary files /dev/null and b/assets/banner.png differ diff --git a/assets/icon.png b/assets/icon.png new file mode 100644 index 0000000..2126498 Binary files /dev/null and b/assets/icon.png differ diff --git a/config.example.json b/config.example.json new file mode 100644 index 0000000..a84dd31 --- /dev/null +++ b/config.example.json @@ -0,0 +1,3 @@ +{ + "OPENAI_API_KEY": "sk-..." +} diff --git a/docker/.dockerignore b/docker/.dockerignore new file mode 100644 index 0000000..474ec69 --- /dev/null +++ b/docker/.dockerignore @@ -0,0 +1,5 @@ +../.git +../.venv +../tests +../docs +../docker diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 0000000..99b64f1 --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,26 @@ +# CPU-only image; Imagent calls the OpenAI API and does no local inference. +FROM python:3.11-slim + +ARG COMFYUI_REF=master + +ENV PYTHONUNBUFFERED=1 \ + PIP_NO_CACHE_DIR=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 + +RUN apt-get update && apt-get install -y --no-install-recommends \ + git build-essential \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /opt +RUN git clone --depth 1 --branch ${COMFYUI_REF} \ + https://github.com/comfyanonymous/ComfyUI.git /opt/ComfyUI + +WORKDIR /opt/ComfyUI +RUN pip install --extra-index-url https://download.pytorch.org/whl/cpu \ + -r requirements.txt + +# Imagent's only runtime dep beyond ComfyUI (Pillow/numpy/torch ship with ComfyUI). +RUN pip install "openai>=1.40.0" + +EXPOSE 8188 +CMD ["python", "main.py", "--listen", "0.0.0.0", "--port", "8188", "--cpu"] diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml new file mode 100644 index 0000000..e1390b0 --- /dev/null +++ b/docker/docker-compose.yml @@ -0,0 +1,26 @@ +services: + comfyui: + build: + context: . + dockerfile: Dockerfile + args: + # Override with `docker compose build --build-arg COMFYUI_REF=v0.3.27` + COMFYUI_REF: master + image: comfyui-imagent-dev:latest + container_name: comfyui-imagent-dev + ports: + - "8188:8188" + environment: + # Checked before config.json by imagent/client.py. Set OPENAI_API_KEY in + # your shell before `up`, or it passes through empty. + OPENAI_API_KEY: ${OPENAI_API_KEY:-} + volumes: + # The node under development: bind-mount the repo into custom_nodes so + # edits reload on container restart. + - ../:/opt/ComfyUI/custom_nodes/comfyui-imagent:rw + # Persistent ComfyUI user dir. + - comfyui_user:/opt/ComfyUI/user + restart: unless-stopped + +volumes: + comfyui_user: diff --git a/imagent/__init__.py b/imagent/__init__.py new file mode 100644 index 0000000..c27ddf0 --- /dev/null +++ b/imagent/__init__.py @@ -0,0 +1,3 @@ +"""imagent package — OpenAI gpt-image generation and editing for ComfyUI.""" + +__version__ = "0.1.0" diff --git a/imagent/build.py b/imagent/build.py new file mode 100644 index 0000000..856617d --- /dev/null +++ b/imagent/build.py @@ -0,0 +1,121 @@ +"""Pure request-building / dispatch logic for the two nodes. + +This module deliberately imports nothing from ComfyUI (`comfy_api`), so the whole +behavior — model gating, kwargs assembly, error handling — is unit-testable +without a running ComfyUI. The IO node wrappers in node_generate.py / +node_edit.py only marshal widget values into these functions. + +Public entry points: run_generate(), run_edit(). +""" +from __future__ import annotations + +import io +import logging + +from PIL import Image + +from . import client as client_mod +from . import image_io +from . import params as params_mod + +log = logging.getLogger("imagent") + +_NO_KEY = "Error: OPENAI_API_KEY not found (set the env var or config.json)." +_NO_IMAGES = "Error: API returned no images." + + +def _named_png(pil_image: Image.Image, name: str) -> io.BytesIO: + buf = io.BytesIO(image_io.pil_to_png_bytes(pil_image)) + buf.name = name # OpenAI SDK infers content-type from the filename + return buf + + +def run_generate(*, prompt, model, size, quality, background, output_format, n, + moderation="auto", custom_width=1024, custom_height=1024, + output_compression=100): + """Text-to-image via images.generate. + + Returns (image_tensor, info_string). Never raises: a failed call yields an + empty image plus a friendly error so a bad call can't crash the queue. + """ + oai = client_mod.get_client() + if oai is None: + log.warning("imagent: %s", _NO_KEY) + return (image_io.empty_image(), _NO_KEY) + try: + resolved_size = params_mod.resolve_size(size, custom_width, custom_height, model) + resolved_background = params_mod.resolve_background(background, model) + kwargs = { + "model": model, "prompt": prompt, "n": n, + "size": resolved_size, "quality": quality, + "background": resolved_background, "output_format": output_format, + } + if output_format in ("jpeg", "webp"): + kwargs["output_compression"] = output_compression + if moderation != "auto": + kwargs["moderation"] = moderation + resp = oai.images.generate(**kwargs) + if not resp.data: + return (image_io.empty_image(), _NO_IMAGES) + info = getattr(resp.data[0], "revised_prompt", None) \ + or f"Generated {len(resp.data)} image(s) with {model}." + return (image_io.batch_from_response_data(resp.data), info) + except Exception as exc: # noqa: BLE001 - node must not crash the queue + log.exception("imagent: image generation failed") + return (image_io.empty_image(), client_mod.friendly_error(exc)) + + +def run_edit(*, prompt, model, size, quality, background, output_format, n, + images, mask=None, input_fidelity="high", moderation="auto", + custom_width=1024, custom_height=1024): + """Image edit / inpaint / multi-reference via images.edit. + + `images` is an iterable of IMAGE tensors (None entries are ignored). + Returns (image_tensor, info_string); never raises (see run_generate). + """ + oai = client_mod.get_client() + if oai is None: + log.warning("imagent: %s", _NO_KEY) + return (image_io.empty_image(), _NO_KEY) + + refs = [t for t in (images or []) if t is not None] + if not refs: + msg = "Error: OpenAIImageEdit requires at least one reference image." + log.warning("imagent: %s", msg) + return (image_io.empty_image(), msg) + + if mask is not None and len(refs) > 1: + msg = ("Error: mask inpainting requires a single reference image " + "(connect only one image).") + log.warning("imagent: %s", msg) + return (image_io.empty_image(), msg) + + try: + resolved_size = params_mod.resolve_size(size, custom_width, custom_height, model) + resolved_background = params_mod.resolve_background(background, model) + caps = params_mod.capabilities_for(model) + # Downscale oversized references to fit the API's input pixel budget. + ref_pils = [image_io.downscale_pil_to_pixel_limit(image_io.tensor_to_pil(t)[0]) + for t in refs] + + kwargs = { + "model": model, "prompt": prompt, + "image": [_named_png(p, f"ref_{i}.png") for i, p in enumerate(ref_pils)], + "n": n, "size": resolved_size, "quality": quality, + "background": resolved_background, "output_format": output_format, + } + if mask is not None: + kwargs["mask"] = image_io.mask_to_named_png(ref_pils[0], mask, "mask.png") + if caps["input_fidelity"]: + kwargs["input_fidelity"] = input_fidelity + if moderation != "auto": + kwargs["moderation"] = moderation + resp = oai.images.edit(**kwargs) + if not resp.data: + return (image_io.empty_image(), _NO_IMAGES) + info = getattr(resp.data[0], "revised_prompt", None) \ + or f"Edited image with {model}." + return (image_io.batch_from_response_data(resp.data), info) + except Exception as exc: # noqa: BLE001 - node must not crash the queue + log.exception("imagent: image edit failed") + return (image_io.empty_image(), client_mod.friendly_error(exc)) diff --git a/imagent/client.py b/imagent/client.py new file mode 100644 index 0000000..6ce6797 --- /dev/null +++ b/imagent/client.py @@ -0,0 +1,76 @@ +"""OpenAI API key resolution, model catalog, and client construction. + +Public entry points: resolve_api_key(), get_client(), friendly_error(); plus the +MODELS list and DEFAULT_MODEL constant. +""" +from __future__ import annotations + +import json +import logging +import os +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from openai import OpenAI + +log = logging.getLogger("imagent") + +# config.json lives at the repo root, next to this package directory. +_CONFIG_PATH = Path(__file__).resolve().parent.parent / "config.json" + +# Current OpenAI image models, verified 2026-06-07. Curated to the three current +# GPT image generations (1, 1.5, 2); dated snapshots and -mini are intentionally +# omitted. DALL-E excluded (shut down 2026-05-12). +MODELS = [ + "gpt-image-2", + "gpt-image-1.5", + "gpt-image-1", +] +DEFAULT_MODEL = "gpt-image-2" + +_PLACEHOLDERS = {"", "sk-...", "your_openai_api_key_here"} + + +def resolve_api_key() -> str | None: + """Return the API key from env (preferred) or config.json, else None.""" + env_key = os.environ.get("OPENAI_API_KEY") + if env_key and env_key not in _PLACEHOLDERS: + return env_key + try: + if _CONFIG_PATH.exists(): + cfg = json.loads(_CONFIG_PATH.read_text()) + file_key = cfg.get("OPENAI_API_KEY", "") + if isinstance(file_key, str) and file_key and file_key not in _PLACEHOLDERS: + return file_key + except (json.JSONDecodeError, OSError) as exc: + log.warning("imagent: could not read %s: %s", _CONFIG_PATH, exc) + return None + + +def get_client() -> "OpenAI | None": + """Construct an OpenAI client, or None if no key is available.""" + key = resolve_api_key() + if not key: + return None + from openai import OpenAI + return OpenAI(api_key=key) + + +def friendly_error(exc: Exception) -> str: + """Turn an OpenAI exception into an actionable message. + + Prefers the exception's HTTP status_code (OpenAI APIError exposes it) and + falls back to anchored substrings, so unrelated text containing 'verif' or + '429' is not misclassified. + """ + msg = str(exc) + lower = msg.lower() + status = getattr(exc, "status_code", None) + is_verif = "organization verification" in lower or "must be verified" in lower + if status == 403 or "error code: 403" in lower or is_verif: + return ("Error: 403 — your OpenAI org may need API Organization " + "Verification before gpt-image models can be used. Details: " + msg) + if status == 429 or "error code: 429" in lower: + return "Error: 429 — rate limit / quota exceeded. Details: " + msg + return f"Error: {msg}" diff --git a/imagent/image_io.py b/imagent/image_io.py new file mode 100644 index 0000000..2426674 --- /dev/null +++ b/imagent/image_io.py @@ -0,0 +1,122 @@ +"""Conversions between ComfyUI tensors, PIL images, masks, and base64 payloads. + +Public helpers: pil_to_tensor, tensor_to_pil, b64_to_tensor, empty_image, +pil_to_png_bytes, mask_to_alpha_png_bytes, mask_to_named_png, +downscale_pil_to_pixel_limit, url_to_tensor, response_item_to_tensor, +batch_from_response_data. +""" +from __future__ import annotations + +import base64 +import io + +import numpy as np +import torch +from PIL import Image + + +def pil_to_tensor(pil_image: Image.Image) -> torch.Tensor: + """PIL image -> ComfyUI IMAGE tensor [1, H, W, 3] float32 in 0..1 (RGB).""" + img = pil_image.convert("RGB") + arr = np.array(img).astype(np.float32) / 255.0 + return torch.from_numpy(arr)[None, ] + + +def tensor_to_pil(image_tensor: torch.Tensor) -> list[Image.Image]: + """ComfyUI IMAGE tensor [B, H, W, C] -> list of PIL images (one per batch item).""" + arr = image_tensor.detach().cpu().numpy() + out: list[Image.Image] = [] + for i in range(arr.shape[0]): + a = (np.clip(arr[i], 0.0, 1.0) * 255).astype(np.uint8) + mode = "RGBA" if a.shape[-1] == 4 else "RGB" + out.append(Image.fromarray(a, mode)) + return out + + +def b64_to_tensor(b64_str: str) -> torch.Tensor: + """Base64 PNG/JPEG/WebP -> [1, H, W, 3] float32 tensor.""" + data = base64.b64decode(b64_str) + pil = Image.open(io.BytesIO(data)) + return pil_to_tensor(pil) + + +def empty_image(height: int = 512, width: int = 512) -> torch.Tensor: + """Black RGB placeholder tensor for error returns.""" + return torch.zeros((1, height, width, 3), dtype=torch.float32) + + +def pil_to_png_bytes(pil_image: Image.Image) -> bytes: + """PIL image -> PNG bytes.""" + buf = io.BytesIO() + pil_image.save(buf, format="PNG") + return buf.getvalue() + + +def mask_to_alpha_png_bytes(source_pil: Image.Image, mask_tensor: torch.Tensor) -> bytes: + """Build an RGBA PNG where the mask-selected region is transparent. + + ComfyUI MASK convention: 1.0 = selected. OpenAI edits the *transparent* + region, so selected (1.0) maps to alpha 0. The mask is resized to the + source image dimensions. + """ + m = mask_tensor + if m.dim() == 3: + m = m[0] + m_np = np.clip(m.detach().cpu().numpy(), 0.0, 1.0) + mask_pil = Image.fromarray((m_np * 255).astype(np.uint8), "L").resize(source_pil.size) + mask_resized = np.array(mask_pil).astype(np.float32) / 255.0 + alpha = ((1.0 - mask_resized) * 255).astype(np.uint8) # selected -> 0 (editable) + rgba = source_pil.convert("RGBA") + rgba.putalpha(Image.fromarray(alpha, "L")) + return pil_to_png_bytes(rgba) + + +def mask_to_named_png(source_pil: Image.Image, mask_tensor: torch.Tensor, name: str) -> io.BytesIO: + """mask_to_alpha_png_bytes wrapped as a named BytesIO for the OpenAI SDK.""" + buf = io.BytesIO(mask_to_alpha_png_bytes(source_pil, mask_tensor)) + buf.name = name + return buf + + +# Default budget for downscaling reference images before upload (2048*2048). +MAX_INPUT_PIXELS = 4_194_304 + + +def downscale_pil_to_pixel_limit(pil_image: Image.Image, + max_pixels: int = MAX_INPUT_PIXELS) -> Image.Image: + """Downscale a PIL image so total pixels <= max_pixels, preserving aspect ratio. + + Returns the original image unchanged when it already fits. Mirrors the + built-in ComfyUI node's downscale_image_tensor robustness. + """ + w, h = pil_image.size + if w * h <= max_pixels: + return pil_image + scale = (max_pixels / float(w * h)) ** 0.5 + return pil_image.resize((max(1, int(w * scale)), max(1, int(h * scale)))) + + +def url_to_tensor(url: str) -> torch.Tensor: + """Download an image URL and convert it to a [1, H, W, 3] tensor.""" + import urllib.request + with urllib.request.urlopen(url, timeout=30) as resp: # noqa: S310 - caller owns URL provenance + data = resp.read() + return pil_to_tensor(Image.open(io.BytesIO(data))) + + +def response_item_to_tensor(item) -> torch.Tensor: + """Convert an OpenAI image response item (b64_json or url) to a tensor.""" + b64 = getattr(item, "b64_json", None) + if b64: + return b64_to_tensor(b64) + url = getattr(item, "url", None) + if url: + return url_to_tensor(url) + raise ValueError("response item has neither b64_json nor url") + + +def batch_from_response_data(data: list) -> torch.Tensor: + """OpenAI response `.data` list -> [B, H, W, 3] batch tensor (b64 or url items).""" + if not data: + raise ValueError("data must contain at least one response item") + return torch.cat([response_item_to_tensor(d) for d in data], dim=0) diff --git a/imagent/node_edit.py b/imagent/node_edit.py new file mode 100644 index 0000000..517f849 --- /dev/null +++ b/imagent/node_edit.py @@ -0,0 +1,88 @@ +"""ComfyUI node: image edit / inpaint via OpenAI images.edit. + +Like node_generate, `model` is a DynamicCombo with a nested size combo. Reference +images use an auto-growing input (up to 16). Behavior lives in build.run_edit(). +""" +from __future__ import annotations + +from comfy_api.latest import IO + +from . import build +from . import params as params_mod +from .node_generate import _gpt_image_2_inputs, unpack_size + +_MAX_REFS = 16 + + +def _legacy_edit_inputs(): + """gpt-image-1 / gpt-image-1.5: base sizes, transparent allowed, input_fidelity.""" + return [ + IO.Combo.Input("size", options=params_mod.SIZES_BASE, default="auto", + tooltip="Output size (base sizes supported by gpt-image-1.x)."), + IO.Combo.Input("background", options=params_mod.BACKGROUNDS, default="auto", + tooltip="Background handling. 'transparent' needs png or webp output."), + IO.Combo.Input("input_fidelity", options=params_mod.INPUT_FIDELITIES, default="high", + tooltip="How closely to preserve the reference image: 'high' " + "(default) or 'low'."), + ] + + +class OpenAIImageEdit(IO.ComfyNode): + @classmethod + def define_schema(cls): + return IO.Schema( + node_id="ImagentOpenAIImageEdit", + display_name="🤖 Imagent: OpenAI Image Edit", + category="Imagent", + description="BYOK image edit / inpaint / multi-reference compose via OpenAI " + "gpt-image models (direct API, no proxy).", + inputs=[ + IO.String.Input("prompt", multiline=True, default="", + placeholder="Describe the edit to apply…", + tooltip="Describe the edit to apply to the reference image(s)."), + IO.DynamicCombo.Input( + "model", + tooltip="OpenAI gpt-image model. Switching this changes the options below.", + options=[ + IO.DynamicCombo.Option("gpt-image-2", _gpt_image_2_inputs()), + IO.DynamicCombo.Option("gpt-image-1.5", _legacy_edit_inputs()), + IO.DynamicCombo.Option("gpt-image-1", _legacy_edit_inputs()), + ], + ), + IO.Autogrow.Input( + "images", + template=IO.Autogrow.TemplateNames( + IO.Image.Input("image"), + names=[f"image_{i}" for i in range(1, _MAX_REFS + 1)], + min=1), + tooltip=f"Reference image(s) to edit; grows up to {_MAX_REFS}. A mask " + "(if used) applies to the first image and requires exactly one image."), + IO.Combo.Input("quality", options=params_mod.QUALITIES, default="auto", + tooltip="Rendering quality. 'auto' lets the model decide; " + "higher quality costs more."), + IO.Combo.Input("output_format", options=params_mod.FORMATS, default="png", + tooltip="Image file format returned by the API."), + IO.Combo.Input("moderation", options=params_mod.MODERATIONS, default="auto", + tooltip="Content moderation strictness. 'auto' (default) or " + "'low'. Sent only when not 'auto'."), + IO.Int.Input("n", default=1, min=params_mod.N_MIN, max=params_mod.N_MAX, step=1, + tooltip="How many images to generate (1-8)."), + IO.Mask.Input("mask", optional=True, + tooltip="Inpaint mask: white marks the region to edit. Requires a " + "single reference image."), + ], + outputs=[IO.Image.Output(display_name="image")], + ) + + @classmethod + def execute(cls, prompt, model, images, quality, output_format, moderation, n, + mask=None) -> IO.NodeOutput: + size, custom_width, custom_height = unpack_size(model) + ref_tensors = [t for t in (images or {}).values() if t is not None] + img, _info = build.run_edit( + prompt=prompt, model=model["model"], size=size, quality=quality, + background=model["background"], output_format=output_format, n=n, + images=ref_tensors, mask=mask, + input_fidelity=model.get("input_fidelity", "high"), moderation=moderation, + custom_width=custom_width, custom_height=custom_height) + return IO.NodeOutput(img) diff --git a/imagent/node_generate.py b/imagent/node_generate.py new file mode 100644 index 0000000..2d1d396 --- /dev/null +++ b/imagent/node_generate.py @@ -0,0 +1,111 @@ +"""ComfyUI node: text-to-image via OpenAI images.generate. + +Declared with the comfy_api IO schema. `model` is a DynamicCombo (switching the +model swaps the size/background widgets). Within gpt-image-2, `size` is itself a +nested DynamicCombo so custom_width/height appear only under size='custom'. All +behavior lives in the import-clean build.run_generate(); this file is a thin +schema + marshaling wrapper (so it requires ComfyUI to import). +""" +from __future__ import annotations + +from comfy_api.latest import IO + +from . import build +from . import params as params_mod + + +def _custom_dim(name: str, axis: str): + return IO.Int.Input(name, default=1024, min=params_mod.CUSTOM_DIM_MIN, + max=params_mod.CUSTOM_DIM_MAX, step=params_mod.CUSTOM_DIM_STEP, + tooltip=f"Used when size = custom. Multiple of 16; with {axis}: " + "aspect <= 3:1 and total pixels 655,360-8,294,400.") + + +def _size_combo_gpt2(): + """gpt-image-2 size as a nested DynamicCombo: width/height appear only for 'custom'.""" + options = [IO.DynamicCombo.Option(s, []) for s in params_mod.SIZES_GPT2 if s != "custom"] + options.append(IO.DynamicCombo.Option( + "custom", [_custom_dim("custom_width", "height"), _custom_dim("custom_height", "width")])) + return IO.DynamicCombo.Input( + "size", tooltip="Output size. Choose 'custom' to set an exact width and height.", + options=options) + + +def _gpt_image_2_inputs(): + """Per-model widgets for gpt-image-2: hi-res/custom sizes, opaque/auto background.""" + return [ + _size_combo_gpt2(), + IO.Combo.Input("background", options=params_mod.BACKGROUNDS_GPT2, default="auto", + tooltip="Background handling. gpt-image-2 does not support transparent."), + ] + + +def _legacy_inputs(): + """Per-model widgets for gpt-image-1 / gpt-image-1.5: base sizes, transparent allowed.""" + return [ + IO.Combo.Input("size", options=params_mod.SIZES_BASE, default="auto", + tooltip="Output size (base sizes supported by gpt-image-1.x)."), + IO.Combo.Input("background", options=params_mod.BACKGROUNDS, default="auto", + tooltip="Background handling. 'transparent' needs png or webp output."), + ] + + +def unpack_size(model: dict): + """Pull (size, custom_width, custom_height) from a model dict. + + For gpt-image-2 `size` is a nested DynamicCombo (dict); for gpt-image-1.x it's + a plain combo (string). + """ + sz = model["size"] + if isinstance(sz, dict): + return sz["size"], sz.get("custom_width", 1024), sz.get("custom_height", 1024) + return sz, 1024, 1024 + + +class OpenAIImageGenerate(IO.ComfyNode): + @classmethod + def define_schema(cls): + return IO.Schema( + node_id="ImagentOpenAIImage", + display_name="🤖 Imagent: OpenAI Image", + category="Imagent", + description="BYOK text-to-image via OpenAI gpt-image models (direct API, no proxy).", + inputs=[ + IO.String.Input("prompt", multiline=True, default="", + placeholder="Describe the image to generate…", + tooltip="Text description of the image to generate."), + IO.DynamicCombo.Input( + "model", + tooltip="OpenAI gpt-image model. Switching this changes the options below.", + options=[ + IO.DynamicCombo.Option("gpt-image-2", _gpt_image_2_inputs()), + IO.DynamicCombo.Option("gpt-image-1.5", _legacy_inputs()), + IO.DynamicCombo.Option("gpt-image-1", _legacy_inputs()), + ], + ), + IO.Combo.Input("quality", options=params_mod.QUALITIES, default="auto", + tooltip="Rendering quality. 'auto' lets the model decide; " + "higher quality costs more."), + IO.Combo.Input("output_format", options=params_mod.FORMATS, default="png", + tooltip="Image file format returned by the API."), + IO.Int.Input("output_compression", default=100, min=0, max=100, + tooltip="Compression 0-100. Applies only to jpeg and webp output."), + IO.Combo.Input("moderation", options=params_mod.MODERATIONS, default="auto", + tooltip="Content moderation strictness. 'auto' (default) or " + "'low' (less restrictive). Sent only when not 'auto'."), + IO.Int.Input("n", default=1, min=params_mod.N_MIN, max=params_mod.N_MAX, step=1, + tooltip="How many images to generate (1-8)."), + ], + outputs=[IO.Image.Output(display_name="image")], + ) + + @classmethod + def execute(cls, prompt, model, quality, output_format, output_compression, + moderation, n) -> IO.NodeOutput: + size, custom_width, custom_height = unpack_size(model) + img, _info = build.run_generate( + prompt=prompt, model=model["model"], size=size, quality=quality, + background=model["background"], output_format=output_format, n=n, + moderation=moderation, custom_width=custom_width, custom_height=custom_height, + output_compression=output_compression) + return IO.NodeOutput(img) diff --git a/imagent/params.py b/imagent/params.py new file mode 100644 index 0000000..ab566e0 --- /dev/null +++ b/imagent/params.py @@ -0,0 +1,116 @@ +"""Parameter constants, per-model capabilities, and size validation. + +Public entry points: capabilities_for(), validate_custom_dimensions(), resolve_size(), +resolve_background(); plus the SIZE_PRESETS / QUALITIES / BACKGROUNDS / FORMATS / +INPUT_FIDELITIES / MODERATIONS lists and the CUSTOM_DIM_* widget bounds. +""" +from __future__ import annotations + +import logging + +log = logging.getLogger("imagent") + +# Size presets. The base sizes work on every gpt-image model; the high-res presets +# and 'custom' are gpt-image-2 only (gated in resolve_size to stay cross-model safe). +_BASE_SIZES = ["auto", "1024x1024", "1024x1536", "1536x1024"] +_GPT2_SIZES = ["2048x2048", "2048x1152", "1152x2048", "3840x2160", "2160x3840"] +_GPT2_ONLY_SIZES = set(_GPT2_SIZES) +SIZE_PRESETS = [*_BASE_SIZES, *_GPT2_SIZES, "custom"] + +# Per-model size option lists for the IO DynamicCombo (what each model can pick). +SIZES_BASE = list(_BASE_SIZES) # gpt-image-1 / 1.5 +SIZES_GPT2 = [*_BASE_SIZES, *_GPT2_SIZES, "custom"] # gpt-image-2 + +# custom_width / custom_height INT-widget bounds (mirror the official node). +CUSTOM_DIM_MIN, CUSTOM_DIM_MAX, CUSTOM_DIM_STEP = 1024, 3840, 16 + +# Per-model background option lists (gpt-image-2 cannot do transparent). +BACKGROUNDS = ["auto", "opaque", "transparent"] +BACKGROUNDS_GPT2 = ["auto", "opaque"] + +QUALITIES = ["auto", "low", "medium", "high"] +FORMATS = ["png", "jpeg", "webp"] +INPUT_FIDELITIES = ["high", "low"] +MODERATIONS = ["auto", "low"] + +# Widget bounds shared by both nodes' IO schemas. +N_MIN, N_MAX = 1, 8 + +# Per-model capability flags. Keys (all default False for unknown models): +# custom_size - arbitrary WxH sizes (gpt-image-2 only) +# input_fidelity - edit input_fidelity param (gpt-image-1.x) +# transparent_background - background="transparent" (NOT gpt-image-2, which errors) +_DEFAULT_CAPS = { + "custom_size": False, "input_fidelity": False, "transparent_background": False, +} + +CAPABILITIES: dict[str, dict[str, bool]] = { + "gpt-image-2": {"custom_size": True, "input_fidelity": False, "transparent_background": False}, + "gpt-image-1.5": {"custom_size": False, "input_fidelity": True, "transparent_background": True}, + "gpt-image-1": {"custom_size": False, "input_fidelity": True, "transparent_background": True}, +} + +# gpt-image-2 custom-resolution constraints (mirror the official node, 2026-06-08). +_MAX_EDGE = 3840 +_MIN_TOTAL_PX = 655_360 +_MAX_TOTAL_PX = 8_294_400 + + +def capabilities_for(model: str) -> dict[str, bool]: + """Return a copy of the capability flags for *model*, defaulting to all-False.""" + return dict(CAPABILITIES.get(model, _DEFAULT_CAPS)) + + +def validate_custom_dimensions(width: int, height: int) -> str: + """Validate gpt-image-2 custom WxH dimensions; return the normalized 'WxH'. + + Rules match the official node: multiples of 16, max edge 3840, aspect <= 3:1, + and total pixels between 655,360 and 8,294,400. + """ + try: + w, h = int(width), int(height) + except (TypeError, ValueError): + raise ValueError(f"custom dimensions must be integers, got {width!r}x{height!r}") from None + if w <= 0 or h <= 0: + raise ValueError(f"custom dimensions must be positive, got {w}x{h}") + if w % 16 or h % 16: + raise ValueError(f"custom width and height must be multiples of 16, got {w}x{h}") + if max(w, h) > _MAX_EDGE: + raise ValueError(f"custom max edge must be <= {_MAX_EDGE}, got {w}x{h}") + if max(w, h) / min(w, h) > 3: + raise ValueError(f"custom aspect ratio must not exceed 3:1, got {w}x{h}") + total = w * h + if not (_MIN_TOTAL_PX <= total <= _MAX_TOTAL_PX): + raise ValueError( + f"custom total pixels must be between {_MIN_TOTAL_PX} and {_MAX_TOTAL_PX}, got {total}") + return f"{w}x{h}" + + +def resolve_size(size: str, custom_width: int, custom_height: int, model: str) -> str: + """Resolve the size parameter to a value the selected model accepts. + + 'custom' and the high-res presets are gpt-image-2 only; for other models they + fall back to 'auto' so switching models never sends an unsupported size. + """ + caps = capabilities_for(model) + if size == "custom": + if not caps["custom_size"]: + log.info("imagent: %s does not support custom sizes; using 'auto'.", model) + return "auto" + return validate_custom_dimensions(custom_width, custom_height) + if size in _GPT2_ONLY_SIZES and not caps["custom_size"]: + log.info("imagent: size %s is gpt-image-2 only; using 'auto' for %s.", size, model) + return "auto" + return size + + +def resolve_background(background: str, model: str) -> str: + """Drop an unsupported transparent background to 'auto'. + + gpt-image-2 errors on background='transparent'; only gpt-image-1.x supports + it. Falling back to 'auto' keeps the node compatible across model switches. + """ + if background == "transparent" and not capabilities_for(model)["transparent_background"]: + log.info("imagent: %s does not support a transparent background; using 'auto'.", model) + return "auto" + return background diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..8d4bc6a --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,50 @@ +[build-system] +requires = ["setuptools>=77"] +build-backend = "setuptools.build_meta" + +[project] +name = "comfyui-imagent" +dynamic = ["version"] +description = "ComfyUI nodes for OpenAI gpt-image generation and editing" +readme = "README.md" +license = "MIT" +requires-python = ">=3.10" +authors = [{ name = "ComfyUI-Imagent contributors" }] +keywords = ["comfyui", "openai", "gpt-image", "image-generation"] +classifiers = [ + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", +] +dependencies = ["openai>=1.40.0", "Pillow>=10.0", "numpy"] + +[project.optional-dependencies] +test = ["pytest>=8.0", "torch"] + +[project.urls] +Repository = "https://github.com/agarzon/ComfyUI-Imagent" + +# ComfyUI Registry metadata (used by `comfy node publish`). +[tool.comfy] +PublisherId = "agarzon" +DisplayName = "Imagent" +requires-comfyui = ">=0.23.0" +Icon = "https://raw.githubusercontent.com/agarzon/ComfyUI-Imagent/main/assets/icon.png" +Banner = "https://raw.githubusercontent.com/agarzon/ComfyUI-Imagent/main/assets/banner.png" + +[tool.comfy.version] +# version is dynamic (imagent.__version__); comfy-cli reads it from this source file +path = "imagent/__init__.py" + +[tool.setuptools.dynamic] +version = { attr = "imagent.__version__" } + +[tool.setuptools.packages.find] +include = ["imagent*"] +exclude = ["tests*", "docker*", "docs*"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = ["."] +addopts = "--import-mode=importlib" +xfail_strict = true diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..51e502e --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +openai>=1.40.0 +Pillow>=10.0 +numpy 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..2fb52f9 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,2 @@ +"""Shared pytest fixtures for ComfyUI-Imagent.""" +from __future__ import annotations diff --git a/tests/test_client.py b/tests/test_client.py new file mode 100644 index 0000000..417a644 --- /dev/null +++ b/tests/test_client.py @@ -0,0 +1,72 @@ +import json +import imagent.client as client + + +def test_env_var_takes_precedence(monkeypatch, tmp_path): + monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env") + monkeypatch.setattr(client, "_CONFIG_PATH", tmp_path / "config.json") + (tmp_path / "config.json").write_text(json.dumps({"OPENAI_API_KEY": "sk-from-file"})) + assert client.resolve_api_key() == "sk-from-env" + + +def test_config_file_fallback(monkeypatch, tmp_path): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + cfg = tmp_path / "config.json" + cfg.write_text(json.dumps({"OPENAI_API_KEY": "sk-from-file"})) + monkeypatch.setattr(client, "_CONFIG_PATH", cfg) + assert client.resolve_api_key() == "sk-from-file" + + +def test_missing_key_returns_none(monkeypatch, tmp_path): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setattr(client, "_CONFIG_PATH", tmp_path / "nope.json") + assert client.resolve_api_key() is None + + +def test_placeholder_key_is_ignored(monkeypatch, tmp_path): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + cfg = tmp_path / "config.json" + cfg.write_text(json.dumps({"OPENAI_API_KEY": "sk-..."})) + monkeypatch.setattr(client, "_CONFIG_PATH", cfg) + assert client.resolve_api_key() is None + + +def test_model_catalog_defaults(): + assert client.DEFAULT_MODEL == "gpt-image-2" + assert set(client.MODELS) == {"gpt-image-1", "gpt-image-1.5", "gpt-image-2"} + assert "gpt-image-1-mini" not in client.MODELS + assert "gpt-image-2-2026-04-21" not in client.MODELS + assert all("dall-e" not in m for m in client.MODELS) + + +def test_friendly_error_maps_403_and_429(): + assert "verif" in client.friendly_error(Exception("Error code: 403 - org not verified")).lower() + assert "rate" in client.friendly_error(Exception("Error code: 429 - too many requests")).lower() + + +def test_malformed_config_returns_none(monkeypatch, tmp_path, caplog): + import logging + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + cfg = tmp_path / "config.json" + cfg.write_text("{bad json") + monkeypatch.setattr(client, "_CONFIG_PATH", cfg) + with caplog.at_level(logging.WARNING, logger="imagent"): + assert client.resolve_api_key() is None + assert "could not read" in caplog.text + + +def test_env_placeholder_is_ignored(monkeypatch, tmp_path): + monkeypatch.setenv("OPENAI_API_KEY", "your_openai_api_key_here") + monkeypatch.setattr(client, "_CONFIG_PATH", tmp_path / "nope.json") + assert client.resolve_api_key() is None + + +def test_friendly_error_uses_status_code(): + class FakeAPIError(Exception): + status_code = 403 + + class FakeRateError(Exception): + status_code = 429 + + assert "verif" in client.friendly_error(FakeAPIError("denied")).lower() + assert "rate" in client.friendly_error(FakeRateError("slow down")).lower() diff --git a/tests/test_image_io.py b/tests/test_image_io.py new file mode 100644 index 0000000..021e37b --- /dev/null +++ b/tests/test_image_io.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import base64 +import io +import types +import numpy as np +import torch +from PIL import Image +from imagent import image_io + + +def _png_b64(color=(10, 20, 30), size=(32, 48)): + img = Image.new("RGB", size, color) + buf = io.BytesIO() + img.save(buf, format="PNG") + return base64.b64encode(buf.getvalue()).decode() + + +def test_pil_to_tensor_shape_and_range(): + t = image_io.pil_to_tensor(Image.new("RGB", (8, 4), (255, 0, 0))) + assert t.shape == (1, 4, 8, 3) # [B, H, W, C] + assert t.dtype == torch.float32 + assert 0.0 <= float(t.min()) and float(t.max()) <= 1.0 + assert abs(float(t[0, 0, 0, 0]) - 1.0) < 1e-6 # red channel + + +def test_b64_to_tensor_roundtrip(): + t = image_io.b64_to_tensor(_png_b64(size=(32, 48))) + assert t.shape == (1, 48, 32, 3) + + +def test_empty_image_is_black_rgb(): + e = image_io.empty_image(64, 80) + assert e.shape == (1, 64, 80, 3) + assert float(e.abs().sum()) == 0.0 + + +def test_tensor_to_pil_returns_one_per_batch(): + t = torch.zeros((2, 4, 8, 3), dtype=torch.float32) + pils = image_io.tensor_to_pil(t) + assert len(pils) == 2 + assert pils[0].size == (8, 4) # PIL size is (W, H) + mid = image_io.tensor_to_pil(torch.full((1, 2, 2, 3), 0.5, dtype=torch.float32))[0] + assert int(np.array(mid)[0, 0, 0]) == 127 # 0.5 * 255 -> 127 + + +def test_mask_selected_region_becomes_transparent(): + source = Image.new("RGB", (4, 4), (255, 255, 255)) + mask = torch.zeros((4, 4), dtype=torch.float32) + mask[0:2, 0:2] = 1.0 + png_bytes = image_io.mask_to_alpha_png_bytes(source, mask) + out = Image.open(io.BytesIO(png_bytes)).convert("RGBA") + alpha = np.array(out)[:, :, 3] + assert alpha[0, 0] == 0 # selected -> transparent (editable) + assert alpha[3, 3] == 255 # unselected -> opaque (preserved) + + +def test_downscale_pil_to_pixel_limit_shrinks_large_image(): + big = Image.new("RGB", (4000, 4000), (5, 5, 5)) + out = image_io.downscale_pil_to_pixel_limit(big, max_pixels=1_000_000) + w, h = out.size + assert w * h <= 1_000_000 + assert abs((w / h) - 1.0) < 0.05 # aspect ratio preserved + + +def test_downscale_pil_to_pixel_limit_keeps_small_image(): + small = Image.new("RGB", (64, 48), (5, 5, 5)) + out = image_io.downscale_pil_to_pixel_limit(small, max_pixels=1_000_000) + assert out.size == (64, 48) # unchanged + + +def test_response_item_to_tensor_handles_b64(): + item = types.SimpleNamespace(b64_json=_png_b64(size=(20, 10)), url=None) + t = image_io.response_item_to_tensor(item) + assert t.shape == (1, 10, 20, 3) + + +def test_response_item_to_tensor_handles_url(monkeypatch): + png = base64.b64decode(_png_b64(size=(12, 6))) + + class _FakeResp: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def read(self): + return png + + monkeypatch.setattr("urllib.request.urlopen", lambda url, **kw: _FakeResp()) + item = types.SimpleNamespace(b64_json=None, url="http://example/x.png") + t = image_io.response_item_to_tensor(item) + assert t.shape == (1, 6, 12, 3) + + +def test_batch_from_response_data_concatenates_b64_items(): + data = [types.SimpleNamespace(b64_json=_png_b64(), url=None) for _ in range(3)] + batch = image_io.batch_from_response_data(data) + assert batch.shape[0] == 3 + + +def test_response_item_to_tensor_raises_without_payload(): + import pytest + item = types.SimpleNamespace(b64_json=None, url=None) + with pytest.raises(ValueError): + image_io.response_item_to_tensor(item) + + +def test_mask_3d_shape_is_handled(): + source = Image.new("RGB", (4, 4), (255, 255, 255)) + mask = torch.zeros((1, 4, 4), dtype=torch.float32) # leading batch dim + mask[0, 0:2, 0:2] = 1.0 + png_bytes = image_io.mask_to_alpha_png_bytes(source, mask) + out = Image.open(io.BytesIO(png_bytes)).convert("RGBA") + alpha = np.array(out)[:, :, 3] + assert alpha[0, 0] == 0 + assert alpha[3, 3] == 255 + + +def test_batch_from_response_data_rejects_empty(): + import pytest + with pytest.raises(ValueError): + image_io.batch_from_response_data([]) + + +def test_mask_to_named_png_returns_named_filelike(): + src = Image.new("RGB", (8, 8), (200, 200, 200)) + mask = torch.zeros((8, 8), dtype=torch.float32) + mask[0:4, :] = 1.0 + f = image_io.mask_to_named_png(src, mask, "mask.png") + assert f.name == "mask.png" + out = Image.open(f).convert("RGBA") + assert out.size == (8, 8) diff --git a/tests/test_nodes.py b/tests/test_nodes.py new file mode 100644 index 0000000..4b725c7 --- /dev/null +++ b/tests/test_nodes.py @@ -0,0 +1,288 @@ +import base64 +import io +import types +from pathlib import Path + +from PIL import Image + +import imagent.client as client +from imagent.build import run_generate, run_edit + + +def _b64(color=(1, 2, 3), size=(16, 16)): + buf = io.BytesIO() + Image.new("RGB", size, color).save(buf, format="PNG") + return base64.b64encode(buf.getvalue()).decode() + + +class FakeImages: + def __init__(self, recorder): + self._rec = recorder + + def generate(self, **kwargs): + self._rec["generate"] = kwargs + item = types.SimpleNamespace(b64_json=_b64(), revised_prompt="a revised prompt") + return types.SimpleNamespace(data=[item for _ in range(kwargs.get("n", 1))]) + + def edit(self, **kwargs): + self._rec["edit"] = kwargs + item = types.SimpleNamespace(b64_json=_b64(), revised_prompt=None) + return types.SimpleNamespace(data=[item for _ in range(kwargs.get("n", 1))]) + + +class FakeClient: + def __init__(self, recorder): + self.images = FakeImages(recorder) + + +def test_generate_missing_key_returns_empty(monkeypatch): + monkeypatch.setattr(client, "get_client", lambda: None) + img, info = run_generate( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="png", n=1) + assert img.shape == (1, 512, 512, 3) + assert info.startswith("Error") + + +def test_generate_sends_expected_params_and_decodes(monkeypatch): + rec = {} + monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec)) + img, info = run_generate( + prompt="a cat", model="gpt-image-2", size="1024x1024", quality="high", + background="opaque", output_format="png", n=2) + sent = rec["generate"] + assert sent["model"] == "gpt-image-2" + assert sent["prompt"] == "a cat" + assert sent["n"] == 2 + assert sent["size"] == "1024x1024" + assert sent["quality"] == "high" + assert sent["background"] == "opaque" + assert img.shape[0] == 2 + assert info == "a revised prompt" + + +def test_generate_drops_transparent_background_for_gpt_image_2(monkeypatch): + rec = {} + monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec)) + run_generate( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="transparent", output_format="png", n=1) + assert rec["generate"]["background"] == "auto" # gpt-image-2 can't do transparent + + +def test_generate_keeps_transparent_background_for_gpt_image_1(monkeypatch): + rec = {} + monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec)) + run_generate( + prompt="x", model="gpt-image-1", size="auto", quality="auto", + background="transparent", output_format="png", n=1) + assert rec["generate"]["background"] == "transparent" + + +def test_generate_adds_compression_only_for_jpeg(monkeypatch): + rec = {} + monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec)) + run_generate( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="jpeg", n=1, output_compression=80) + assert rec["generate"]["output_compression"] == 80 + rec.clear() + run_generate( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="png", n=1, output_compression=80) + assert "output_compression" not in rec["generate"] + + +def test_generate_moderation_only_when_not_auto(monkeypatch): + rec = {} + monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec)) + run_generate( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="png", n=1, moderation="low") + assert rec["generate"]["moderation"] == "low" + rec.clear() + run_generate( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="png", n=1, moderation="auto") + assert "moderation" not in rec["generate"] + + +def test_n_max_is_8(): + from imagent import params + assert params.N_MAX == 8 + + +def test_generate_api_error_returns_empty(monkeypatch): + monkeypatch.setattr(client, "get_client", lambda: FakeClient({})) + monkeypatch.setattr( + FakeImages, "generate", + lambda self, **kw: (_ for _ in ()).throw(RuntimeError("network error"))) + img, info = run_generate( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="png", n=1) + assert img.shape == (1, 512, 512, 3) + assert info.startswith("Error") + + +def test_generate_invalid_custom_size_returns_error(monkeypatch): + monkeypatch.setattr(client, "get_client", lambda: FakeClient({})) + img, info = run_generate( + prompt="x", model="gpt-image-2", size="custom", quality="auto", + background="auto", output_format="png", n=1, + custom_width=1000, custom_height=1000) # 1000 is not a multiple of 16 + assert img.shape == (1, 512, 512, 3) + assert info.startswith("Error") + + +def test_generate_custom_dimensions_for_gpt_image_2(monkeypatch): + rec = {} + monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec)) + run_generate( + prompt="x", model="gpt-image-2", size="custom", quality="auto", + background="auto", output_format="png", n=1, + custom_width=2048, custom_height=1152) + assert rec["generate"]["size"] == "2048x1152" + + +def test_generate_hires_preset_gated_to_gpt_image_2(monkeypatch): + rec = {} + monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec)) + run_generate( + prompt="x", model="gpt-image-1", size="2048x2048", quality="auto", + background="auto", output_format="png", n=1) + assert rec["generate"]["size"] == "auto" # gpt-image-1 can't do hi-res + rec.clear() + run_generate( + prompt="x", model="gpt-image-2", size="2048x2048", quality="auto", + background="auto", output_format="png", n=1) + assert rec["generate"]["size"] == "2048x2048" + + +# --------------------------------------------------------------------------- +# OpenAIImageEdit tests +# --------------------------------------------------------------------------- +import torch + + +def _img_tensor(h=16, w=16): + return torch.full((1, h, w, 3), 0.5, dtype=torch.float32) + + +def test_edit_sends_image_list_and_decodes(monkeypatch): + rec = {} + monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec)) + img, info = run_edit( + prompt="make it snowy", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="png", n=1, + images=[_img_tensor(), _img_tensor()]) + sent = rec["edit"] + assert sent["model"] == "gpt-image-2" + assert isinstance(sent["image"], list) and len(sent["image"]) == 2 + assert "mask" not in sent + assert img.shape[0] == 1 + + +def test_edit_includes_mask_when_provided(monkeypatch): + rec = {} + monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec)) + mask = torch.zeros((16, 16), dtype=torch.float32) + mask[0:8, 0:8] = 1.0 + run_edit( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="png", n=1, + images=[_img_tensor()], mask=mask) + assert "mask" in rec["edit"] + + +def test_edit_input_fidelity_only_for_1x(monkeypatch): + rec = {} + monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec)) + run_edit( + prompt="x", model="gpt-image-1", size="auto", quality="auto", + background="auto", output_format="png", n=1, + images=[_img_tensor()], input_fidelity="high") + assert rec["edit"]["input_fidelity"] == "high" + rec.clear() + run_edit( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="png", n=1, + images=[_img_tensor()], input_fidelity="high") + assert "input_fidelity" not in rec["edit"] + + +def test_edit_moderation_only_when_not_auto(monkeypatch): + rec = {} + monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec)) + run_edit( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="png", n=1, + images=[_img_tensor()], moderation="low") + assert rec["edit"]["moderation"] == "low" + rec.clear() + run_edit( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="png", n=1, + images=[_img_tensor()], moderation="auto") + assert "moderation" not in rec["edit"] + + +def test_edit_downscales_large_reference(monkeypatch): + from imagent import image_io + rec = {} + monkeypatch.setattr(client, "get_client", lambda: FakeClient(rec)) + big = torch.full((1, 2100, 2100, 3), 0.5, dtype=torch.float32) + run_edit( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="png", n=1, images=[big]) + sent_img = rec["edit"]["image"][0] + w, h = Image.open(sent_img).size + assert w * h <= image_io.MAX_INPUT_PIXELS + + +def test_edit_missing_image_returns_error(monkeypatch): + monkeypatch.setattr(client, "get_client", lambda: FakeClient({})) + img, info = run_edit( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="png", n=1, images=[]) + assert info.startswith("Error") + assert img.shape == (1, 512, 512, 3) + + +def test_edit_api_error_returns_empty(monkeypatch): + monkeypatch.setattr(client, "get_client", lambda: FakeClient({})) + monkeypatch.setattr( + FakeImages, "edit", + lambda self, **kw: (_ for _ in ()).throw(RuntimeError("network error"))) + img, info = run_edit( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="png", n=1, images=[_img_tensor()]) + assert img.shape == (1, 512, 512, 3) + assert info.startswith("Error") + + +def test_edit_mask_with_multiple_refs_returns_error(monkeypatch): + monkeypatch.setattr(client, "get_client", lambda: FakeClient({})) + mask = torch.zeros((16, 16), dtype=torch.float32) + img, info = run_edit( + prompt="x", model="gpt-image-2", size="auto", quality="auto", + background="auto", output_format="png", n=1, + images=[_img_tensor(), _img_tensor()], mask=mask) + assert img.shape == (1, 512, 512, 3) + assert info.startswith("Error") + + +def test_extension_registers_both_nodes(): + # Registration uses the comfy_api V3 entrypoint, so this needs ComfyUI present. + import pytest + pytest.importorskip("comfy_api") + import asyncio + import importlib.util + root = Path(__file__).resolve().parent.parent + spec = importlib.util.spec_from_file_location("imagent_root_init", root / "__init__.py") + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + nodes = asyncio.run(mod.ImagentExtension().get_node_list()) + schemas = {n.define_schema().node_id: n.define_schema().display_name for n in nodes} + assert set(schemas) == {"ImagentOpenAIImage", "ImagentOpenAIImageEdit"} + assert schemas["ImagentOpenAIImage"] == "🤖 Imagent: OpenAI Image" + assert schemas["ImagentOpenAIImageEdit"] == "🤖 Imagent: OpenAI Image Edit" diff --git a/tests/test_params.py b/tests/test_params.py new file mode 100644 index 0000000..7d586d0 --- /dev/null +++ b/tests/test_params.py @@ -0,0 +1,77 @@ +import pytest +from imagent import params + + +def test_constants(): + assert params.SIZE_PRESETS[0] == "auto" + assert "custom" in params.SIZE_PRESETS + assert params.SIZE_PRESETS[-1] == "custom" + for base in ("1024x1024", "1024x1536", "1536x1024"): + assert base in params.SIZE_PRESETS + for hires in ("2048x2048", "3840x2160", "2160x3840"): + assert hires in params.SIZE_PRESETS + assert (params.CUSTOM_DIM_MIN, params.CUSTOM_DIM_MAX, params.CUSTOM_DIM_STEP) == (1024, 3840, 16) + assert params.QUALITIES == ["auto", "low", "medium", "high"] + assert params.BACKGROUNDS == ["auto", "opaque", "transparent"] + assert params.FORMATS == ["png", "jpeg", "webp"] + assert params.MODERATIONS == ["auto", "low"] + + +def test_capabilities_for_known_and_unknown(): + assert params.capabilities_for("gpt-image-2")["custom_size"] is True + assert params.capabilities_for("gpt-image-2")["input_fidelity"] is False + assert params.capabilities_for("gpt-image-2")["transparent_background"] is False + assert params.capabilities_for("gpt-image-1")["input_fidelity"] is True + assert params.capabilities_for("gpt-image-1")["custom_size"] is False + assert params.capabilities_for("gpt-image-1")["transparent_background"] is True + assert params.capabilities_for("gpt-image-1.5")["transparent_background"] is True + assert params.capabilities_for("mystery") == { + "custom_size": False, "input_fidelity": False, "transparent_background": False} + + +def test_resolve_background_drops_transparent_for_gpt_image_2(): + assert params.resolve_background("transparent", "gpt-image-2") == "auto" + assert params.resolve_background("opaque", "gpt-image-2") == "opaque" + assert params.resolve_background("auto", "gpt-image-2") == "auto" + + +def test_resolve_background_keeps_transparent_for_gpt_image_1x(): + assert params.resolve_background("transparent", "gpt-image-1") == "transparent" + assert params.resolve_background("transparent", "gpt-image-1.5") == "transparent" + + +def test_validate_custom_dimensions_ok(): + assert params.validate_custom_dimensions(1536, 864) == "1536x864" + assert params.validate_custom_dimensions(2048, 1152) == "2048x1152" + + +@pytest.mark.parametrize("w,h", [ + (1000, 1000), # not multiples of 16 + (100, 100), # not multiples of 16 (and too small) + (-16, -16), # negative + (1024, 0), # zero + (1024, 512), # total 524,288 < 655,360 minimum + (3840, 2176), # total 8,355,840 > 8,294,400 maximum + (3840, 1024), # aspect 3.75:1 > 3:1 + (4096, 2160), # max edge 4096 > 3840 +]) +def test_validate_custom_dimensions_rejects(w, h): + with pytest.raises(ValueError): + params.validate_custom_dimensions(w, h) + + +def test_resolve_size_passthrough_for_base_presets(): + assert params.resolve_size("1024x1024", 1024, 1024, "gpt-image-2") == "1024x1024" + assert params.resolve_size("auto", 1024, 1024, "gpt-image-1") == "auto" + + +def test_resolve_size_custom_requires_capability(): + assert params.resolve_size("custom", 1536, 864, "gpt-image-2") == "1536x864" + assert params.resolve_size("custom", 1536, 864, "gpt-image-1") == "auto" + + +def test_resolve_size_hires_presets_gated_to_gpt_image_2(): + assert params.resolve_size("2048x2048", 1024, 1024, "gpt-image-2") == "2048x2048" + assert params.resolve_size("3840x2160", 1024, 1024, "gpt-image-2") == "3840x2160" + assert params.resolve_size("2048x2048", 1024, 1024, "gpt-image-1") == "auto" + assert params.resolve_size("3840x2160", 1024, 1024, "gpt-image-1.5") == "auto" diff --git a/tests/test_tooltips.py b/tests/test_tooltips.py new file mode 100644 index 0000000..778a004 --- /dev/null +++ b/tests/test_tooltips.py @@ -0,0 +1,30 @@ +"""Every node input (including per-model DynamicCombo sub-inputs) carries a tooltip. + +The nodes are declared with the comfy_api IO schema, so this test needs ComfyUI +present; it is skipped when comfy_api is unavailable (e.g. the bare dev venv) and +runs inside the ComfyUI container during smoke testing. +""" +from __future__ import annotations + +import pytest + +pytest.importorskip("comfy_api") + +from imagent.node_edit import OpenAIImageEdit # noqa: E402 +from imagent.node_generate import OpenAIImageGenerate # noqa: E402 + + +def _walk(inp, where): + """Assert this input has a tooltip, recursing into (nested) DynamicCombo options.""" + assert getattr(inp, "tooltip", None) and inp.tooltip.strip(), f"{where}.{inp.id} has no tooltip" + opts = getattr(inp, "options", None) + if opts and hasattr(opts[0], "inputs"): # DynamicCombo (possibly nested) + for opt in opts: + for sub in opt.inputs: + _walk(sub, f"{where}.{inp.id}[{opt.key}]") + + +@pytest.mark.parametrize("node", [OpenAIImageGenerate, OpenAIImageEdit]) +def test_every_input_has_a_tooltip(node): + for inp in node.define_schema().inputs: + _walk(inp, node.__name__)