Add Gemini image generation node and multi-image Omni Flash references
- New APIImageGenerate node (external_api/image): text-to-image and image-to-image via Gemini's Imagen (text-only) and Nano Banana / Nano Banana Pro / 2 (text and multi-reference image-to-image) models. - Omni Flash video now accepts a batch of reference images (not just one), sent in order as separate references in the same call. - Fix Veo/Omni resolution and duration handling: correct 720/1080/4K resolution options per model tier, drop the non-existent veo-3.1-fast-generate-001 id, note Google's 8s lock on 1080/4K. - Fix RunwayML model list: gen4_image/gen4_image_turbo are image models, not video, so drop them from the video node's provider list (tracked as follow-up work to add them to the image node). - Update README and tests accordingly. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014Bwg5aRmgChpwrVnpNgoQm
This commit is contained in:
Andy Ratsirarson
co-authored by
Claude Sonnet 5
Claude Sonnet 5
parent
3ea918e84c
commit
c24e649d1c
@@ -1,8 +1,8 @@
|
||||
# ComfyUI_ExternalAPI
|
||||
|
||||
BYOK (bring-your-own-key) ComfyUI nodes for generating video through external providers — Gemini/Veo, OpenAI Sora, Azure, and RunwayML — via [LiteLLM](https://github.com/BerriAI/litellm). No local GPU-hungry model required; you supply your own provider API key and pay that provider directly.
|
||||
BYOK (bring-your-own-key) ComfyUI nodes for generating video and images through external providers — Gemini (Veo, Omni Flash, Imagen, Nano Banana), OpenAI Sora, Azure, and RunwayML — via [LiteLLM](https://github.com/BerriAI/litellm). No local GPU-hungry model required; you supply your own provider API key and pay that provider directly.
|
||||
|
||||
> **Status:** text-to-video works for all providers; image-to-video (a reference image guiding the animation) works for Gemini/Veo only. Text-to-image is planned as a separate node — see [Roadmap](#roadmap).
|
||||
> **Status:** text-to-video works for all providers; image-to-video (a reference image guiding the animation) works for Gemini (Veo and Omni Flash) only. A separate image-generation node exists for Gemini only (Imagen for text-to-image, Nano Banana for both text-to-image and image-to-image) — see [Roadmap](#roadmap) for the other providers.
|
||||
|
||||
## Features
|
||||
|
||||
@@ -14,10 +14,10 @@ BYOK (bring-your-own-key) ComfyUI nodes for generating video through external pr
|
||||
|
||||
| Provider | `provider` value | Models | Aspect ratios | Resolutions | Duration |
|
||||
|---|---|---|---|---|---|
|
||||
| Google Gemini (Veo + Omni Flash) | `gemini` | Fetched live from your account via `google-genai`; falls back to a built-in list of Veo models (`veo-3.1-fast-generate-preview`, `veo-3.1-fast-generate-001`, `veo-3.1-generate-preview`, `veo-3.1-generate-001`, `veo-3.1-lite-generate-preview`, `veo-2.0-generate-001`) plus Omni Flash models (`gemini-omni-1.1-flash`, `gemini-omni-flash-preview`) if that fails | Veo: 16:9, 9:16 · Omni: 16:9, 9:16 | Veo: 1K, 2K · Omni: 360p, 720p, 1080p, 4k | Veo: 4–8s · Omni: 3–10s |
|
||||
| Google Gemini (Veo + Omni Flash) | `gemini` | Fetched live from your account via `google-genai` (filtered to models Google's own listing says still support video generation); falls back to a built-in list of Veo models (`veo-3.1-fast-generate-preview`, `veo-3.1-generate-preview`, `veo-3.1-generate-001`, `veo-3.1-lite-generate-preview`, `veo-2.0-generate-001`) plus Omni Flash models (`gemini-omni-1.1-flash`, `gemini-omni-flash-preview`) if that fails | Veo: 16:9, 9:16 · Omni: 16:9, 9:16 | Veo: 720, 1080, plus 4K on Veo 3.1 (non-Lite) — Veo 2.0 is 720-only · Omni: 360, 720, 1080, 4K | Veo: 4–8s (1080/4K are locked to 8s by Google, not enforced by this dropdown) · Omni: 3–10s |
|
||||
| OpenAI (Sora) | `openai` | `sora-2`, `sora-2-pro`, `sora-2-pro-high-res` | 16:9, 9:16 | 1K, 2K | 4–12s* |
|
||||
| Azure OpenAI (Sora) | `azure` | `sora-2`, `sora-2-pro`, `sora-2-pro-high-res` | 16:9, 9:16 | 1K, 2K | 4–12s* |
|
||||
| RunwayML | `runwayml` | `gen3a_turbo`, `gen4_turbo`, `gen4_aleph`, `gen4_image`, `gen4_image_turbo` | 16:9, 9:16 | 1K, 2K | 5s or 10s |
|
||||
| RunwayML | `runwayml` | `gen3a_turbo`, `gen4_turbo`, `gen4_aleph` | 16:9, 9:16 | 1K, 2K | 5s or 10s |
|
||||
|
||||
\* OpenAI/Azure's upper duration bound isn't published by the provider; 12s is a conservative default rather than a confirmed hard limit.
|
||||
|
||||
@@ -25,7 +25,20 @@ Gemini's model list is fetched live from `google.genai.Client().models.list()` a
|
||||
|
||||
**Omni Flash is a different API shape than Veo, under the hood.** It runs through litellm's request/response Interactions API (`litellm.interactions.acreate`) instead of the polling `avideo_generation`/`avideo_status`/`avideo_content` calls Veo and the other providers use, so there's no in-progress polling for it — the node just waits on the single call. Only one-shot text-to-video is wired up today; Omni Flash's headline features — image/video-conditioned generation and multi-turn conversational editing — aren't implemented (see [Roadmap](#roadmap)), so its models don't appear in the `image`-source `provider` dropdown. Very large outputs that Google delivers by URI instead of inline base64 aren't downloaded yet and will raise a clear error rather than silently failing.
|
||||
|
||||
**Image-to-video** (a reference image conditioning the result) is currently supported for **Gemini/Veo only**. The other providers are text-to-video only for now — see [Roadmap](#roadmap). The models, aspect ratios, resolutions, and durations in the table above apply to both text- and image-to-video.
|
||||
**Image-to-video** (a reference image conditioning the result) is currently supported for **Gemini (Veo and Omni Flash) only**. The other providers are text-to-video only for now — see [Roadmap](#roadmap). The models, aspect ratios, resolutions, and durations in the table above apply to both text- and image-to-video.
|
||||
|
||||
## Supported image providers & models
|
||||
|
||||
A separate node, **API Image Generate (BYOK)**, handles still images. Gemini only for now:
|
||||
|
||||
| Model family | Example model ids | Text-to-image | Image-to-image | Sizes |
|
||||
|---|---|---|---|---|
|
||||
| Imagen | `imagen-4.0-generate-001`, `imagen-4.0-fast-generate-001`, `imagen-3.0-generate-002` | ✅ | ❌ (no edit endpoint) | 1024x1024, 1792x1024, 1024x1792 |
|
||||
| "Nano Banana" (Gemini's multimodal image models) | `gemini-2.5-flash-image` (Nano Banana), `gemini-3-pro-image-preview` (Nano Banana Pro), `gemini-3.1-flash-image-preview` (Nano Banana 2) | ✅ | ✅, including **multiple** reference images | 1024x1024, 1792x1024, 1024x1792 |
|
||||
|
||||
Imagen calls litellm's `aimage_generation` (Google's `:predict` endpoint) and can only generate from text — it has no image-editing capability in litellm at all, so Imagen models don't appear in the node's `image`-source dropdown. Nano Banana models use the same call for text-to-image, and litellm's `aimage_edit` (Google's `:generateContent` endpoint) for image-to-image — which natively accepts a list of images, so batching multiple images into `reference_image` (the same convention as Omni Flash video, above) sends all of them as references in one call.
|
||||
|
||||
Both model families always return images inline (base64), never a hosted URL — this is fetched from Google directly, no local model.
|
||||
|
||||
## Requirements
|
||||
|
||||
@@ -76,27 +89,41 @@ The node also has an optional `api_key` input that overrides the environment var
|
||||
|
||||
Optional inputs: `api_key` (see [Security note](#security-note)), `poll_interval` (seconds between status checks, default 10), `timeout` (max seconds to wait, default 600).
|
||||
|
||||
### Image-to-video (Gemini/Veo)
|
||||
### Image-to-video (Gemini: Veo & Omni Flash)
|
||||
|
||||
To animate a reference image instead of generating from text alone:
|
||||
|
||||
1. Set `source` to `image`. The node reveals a `provider` dropdown that currently lists only `gemini`.
|
||||
2. Pick `gemini`, then pick a Veo `model` — same models as text-to-video.
|
||||
3. The node reveals: `size` (aspect ratio: 16:9 or 9:16), `resolution` (1K/2K), `seconds` (4–8s), `reference_image`, `person_generation`, and an optional `negative_prompt`.
|
||||
2. Pick `gemini`, then pick either a Veo or an Omni Flash `model` — the same combined list as text-to-video.
|
||||
3. The node reveals settings for that model: both show `size`, `resolution`, `seconds`, and `reference_image`; Veo additionally shows `person_generation` and an optional `negative_prompt` (Omni doesn't have either — it only takes text + image + video content).
|
||||
4. Connect `reference_image` — an `IMAGE` socket, same as any other ComfyUI node that takes an image. Wire in a **Load Image** node to upload/pick a file, or connect any other node's `IMAGE` output directly (including an image generated earlier in the same workflow — no need to save it to disk first). This input is technically optional at the graph level, but image-to-video will fail at runtime with a clear error if left unconnected.
|
||||
5. Enter a `prompt`. It's still required with image-to-video — the prompt directs how the reference image is animated.
|
||||
6. Run the graph. As with text-to-video, the node polls until the video finishes and outputs a `VIDEO`.
|
||||
5. **For multiple reference images (Omni Flash only):** feed a *batch* into `reference_image` instead of a single image — e.g. chain several images through a **Batch Images** node upstream. Omni Flash sends each image in the batch, in order, as its own reference alongside the prompt. Veo does not support this: if it receives a batch of more than one image it raises a clear error rather than silently using just the first (litellm's Veo integration only has a single-image field — see [Roadmap](#roadmap)).
|
||||
6. Enter a `prompt`. It's still required with image-to-video — the prompt directs how the reference image(s) are animated.
|
||||
7. Run the graph. Veo polls until the video finishes; Omni Flash is a single request with no status-polling endpoint, so its progress bar just ticks up with elapsed time rather than tracking real progress — both output a `VIDEO`.
|
||||
|
||||
**Note:** for image-to-video, `person_generation` only offers `allow_adult` (text-to-video also offers `allow_all`). Google restricts generation to adults when a reference photo is provided, so the broader option isn't available here.
|
||||
**Note:** for Veo image-to-video, `person_generation` only offers `allow_adult` (text-to-video also offers `allow_all`). Google restricts generation to adults when a reference photo is provided, so the broader option isn't available here.
|
||||
|
||||
This currently supports exactly **one** reference image per generation. It does not (yet) support multiple reference images, tagging/naming images for reference within the prompt text, or interpolation/video-extension — see [Roadmap](#roadmap).
|
||||
Veo is single-image only; Omni Flash accepts multiple, per above. Neither supports tagging/naming individual images for reference within the prompt text (Omni Flash images are just referenced positionally/by natural language, e.g. "the first image" — there's no `@name` syntax on Google's side), video input (Omni Flash also accepts a short video, not just images), or Veo's interpolation/video-extension — see [Roadmap](#roadmap).
|
||||
|
||||
### Image generation (Gemini: Imagen & Nano Banana)
|
||||
|
||||
1. Add the node: right-click canvas → Add Node → `external_api/image` → **API Image Generate (BYOK)**.
|
||||
2. Enter a text `prompt`.
|
||||
3. Under `source`, leave it on `text` for text-to-image, or set it to `image` for image-to-image (see below).
|
||||
4. Pick `provider` → `gemini`, then a `model` — Imagen or Nano Banana for `text`; Nano Banana only for `image` (Imagen has no edit capability).
|
||||
5. The node reveals `size` (1024x1024 / 1792x1024 / 1024x1792) and `num_images` (1–4); the `image` source also reveals `reference_image`.
|
||||
6. For image-to-image, connect `reference_image` the same way as the video node's — a single image, or a **batch** (via a **Batch Images** node) for multiple references, all sent to the same Nano Banana call.
|
||||
7. Run the graph. This is a single request with no polling (image generation is fast) — it outputs an `IMAGE`, connect it to any node that takes one (e.g. **Save Image**, **Preview Image**, or straight into another generation node).
|
||||
|
||||
Optional input: `api_key` (see [Security note](#security-note)).
|
||||
|
||||
## Roadmap
|
||||
|
||||
- **Image generation** — a separate `APIImageGenerate` node, tracked in [issue #2](https://github.com/andy-ratsirarson/ComfyUI_ExternalAPI/issues/2).
|
||||
- **Image-to-video** — implemented for Gemini/Veo (set `source` to `image`), single reference image only. Support for the other providers (OpenAI/Azure Sora, RunwayML) is future work.
|
||||
- **Multi-image references** — RunwayML supports referencing several images at once (with in-prompt tagging), and Gemini/Veo 3.1 separately supports up to 3 reference images plus last-frame interpolation and video extension. None of this richer, multi-image capability is implemented yet — today's image-to-video is single-image only for both providers that support it at all.
|
||||
- **Gemini Omni Flash: image/video-conditioned generation and conversational editing** — Omni Flash accepts image and video inputs and supports iteratively refining a result across turns via `previous_interaction_id`. Only one-shot text-to-video is wired up today (`gemini_omni` provider, `source: text`); feeding it a reference image/video and chaining edits is not yet implemented.
|
||||
- **Image generation for other providers** — implemented for Gemini only (Imagen + Nano Banana; see above). OpenAI (DALL-E 3, gpt-image-1), Azure, and RunwayML (`gen4_image`/`gen4_image_turbo`, currently unused since they were miscategorized as video models — see below) are future work.
|
||||
- **RunwayML `gen4_image`/`gen4_image_turbo` miscategorization** — these are image models, not video, but were previously listed (and still callable) under the video node's RunwayML provider; they've been removed from there but not yet added to the image node.
|
||||
- **Image-to-video** — implemented for Gemini (both Veo and Omni Flash; set `source` to `image`). Support for the other providers (OpenAI/Azure Sora, RunwayML) is future work.
|
||||
- **Veo multi-image references** — Veo 3.1 itself supports up to 3 reference images plus last-frame interpolation and video extension, but litellm's Veo integration (`avideo_generation`) only exposes a single `image` field with no way to pass more — reaching Veo's real multi-image capability would mean bypassing litellm with a raw HTTP call, which isn't done here. Omni Flash, by contrast, already accepts a batch of any size (see above) since its `input` is just a list we build ourselves.
|
||||
- **Gemini Omni Flash: video input and conversational editing** — Omni Flash also accepts a short video as conditioning input (not just images) and supports iteratively refining a result across turns via `previous_interaction_id`. Neither is wired up yet.
|
||||
- **Gemini Omni Flash: large (uri-delivered) outputs** — Google delivers big videos by URI instead of inline base64; downloading those isn't implemented yet, so such a response raises a clear error instead of returning a video.
|
||||
|
||||
## Development
|
||||
|
||||
+124
-16
@@ -76,8 +76,20 @@ class Model:
|
||||
@staticmethod
|
||||
def image_tensor_to_bytesio(image_tensor):
|
||||
"""Convert a ComfyUI IMAGE tensor ([B, H, W, C] float32 in [0, 1]) into
|
||||
a PNG-encoded BytesIO, for providers whose API needs image bytes/a
|
||||
file-like object (e.g. litellm's avideo_generation input_reference).
|
||||
a single PNG-encoded BytesIO from its first frame, for providers whose
|
||||
API takes exactly one reference image (e.g. litellm's avideo_generation
|
||||
input_reference — Veo only). If the tensor is a batch of more than one
|
||||
image, the rest are silently dropped; callers that accept multiple
|
||||
images should use image_tensor_to_bytesio_list instead."""
|
||||
return Model.image_tensor_to_bytesio_list(image_tensor)[0]
|
||||
|
||||
@staticmethod
|
||||
def image_tensor_to_bytesio_list(image_tensor):
|
||||
"""Convert a ComfyUI IMAGE tensor ([B, H, W, C] float32 in [0, 1]) into
|
||||
one PNG-encoded BytesIO per frame in its batch, for providers that
|
||||
accept multiple reference images. A user builds that batch upstream
|
||||
with ComfyUI's own batching nodes (e.g. Batch Images) — this node has
|
||||
no separate multi-image socket.
|
||||
Imports are deferred: torch/numpy/PIL are only available inside a real
|
||||
ComfyUI install, not in this package's own test environment."""
|
||||
from io import BytesIO
|
||||
@@ -85,12 +97,15 @@ class Model:
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
frame = image_tensor[0] if image_tensor.ndim == 4 else image_tensor
|
||||
array = (frame.cpu().numpy() * 255).astype(np.uint8)
|
||||
buffer = BytesIO()
|
||||
Image.fromarray(array).save(buffer, format="PNG")
|
||||
buffer.seek(0)
|
||||
return buffer
|
||||
frames = image_tensor if image_tensor.ndim == 4 else image_tensor.unsqueeze(0)
|
||||
buffers = []
|
||||
for frame in frames:
|
||||
array = (frame.cpu().numpy() * 255).astype(np.uint8)
|
||||
buffer = BytesIO()
|
||||
Image.fromarray(array).save(buffer, format="PNG")
|
||||
buffer.seek(0)
|
||||
buffers.append(buffer)
|
||||
return buffers
|
||||
|
||||
@classmethod
|
||||
def model_combo_option(cls, mode="text"):
|
||||
@@ -116,14 +131,107 @@ class Model:
|
||||
],
|
||||
)
|
||||
|
||||
def image_to_image(self, *args, **kwargs):
|
||||
raise NotImplementedError("image_to_image not implemented for this model.")
|
||||
# --- Image generation (litellm aimage_generation/aimage_edit) ---
|
||||
# Mirrors the video methods above: list_image_models()/image_settings_inputs()/
|
||||
# image_kwargs() for text-to-image, image_to_image_model_ids()/
|
||||
# image_to_image_settings_inputs()/image_to_image_kwargs() for image-to-image.
|
||||
|
||||
def text_to_image(self, *args, **kwargs):
|
||||
raise NotImplementedError("text_to_image not implemented for this model.")
|
||||
@classmethod
|
||||
def list_image_models(cls) -> list[str]:
|
||||
"""Return this provider's available model ids for text-to-image."""
|
||||
raise NotImplementedError(f"list_image_models not implemented for {cls.__name__}.")
|
||||
|
||||
def text_to_video(self, *args, **kwargs):
|
||||
raise NotImplementedError("text_to_video not implemented for this model.")
|
||||
@classmethod
|
||||
def image_to_image_model_ids(cls) -> list[str]:
|
||||
"""Model ids from list_image_models() that also support image-to-image
|
||||
(a reference image + prompt in, an edited/composed image out). Defaults
|
||||
to all of them; override when only some do."""
|
||||
return cls.list_image_models()
|
||||
|
||||
def text_to_text(self, *args, **kwargs):
|
||||
raise NotImplementedError("text_to_text not implemented for this model.")
|
||||
@classmethod
|
||||
def image_settings_inputs(cls, model_id: str | None = None):
|
||||
"""Build the io.Input list revealed once a specific image model is selected."""
|
||||
raise NotImplementedError(f"{cls.__name__} does not support image generation.")
|
||||
|
||||
@classmethod
|
||||
def image_kwargs(cls, model_id: str, settings: dict) -> dict:
|
||||
"""Translate the generic settings dict into the litellm aimage_generation
|
||||
kwargs this provider's API expects."""
|
||||
raise NotImplementedError(f"image_kwargs not implemented for {cls.__name__}.")
|
||||
|
||||
@classmethod
|
||||
def image_to_image_settings_inputs(cls, model_id: str | None = None):
|
||||
"""Build the io.Input list revealed once a specific model is selected,
|
||||
for the image-to-image ('image' source) path."""
|
||||
raise NotImplementedError(f"{cls.__name__} does not support image-to-image.")
|
||||
|
||||
@classmethod
|
||||
def image_to_image_kwargs(cls, model_id: str, settings: dict) -> dict:
|
||||
"""Translate the generic settings dict (including the reference image
|
||||
tensor) into the litellm aimage_edit kwargs this provider's API expects."""
|
||||
raise NotImplementedError(f"image_to_image_kwargs not implemented for {cls.__name__}.")
|
||||
|
||||
@classmethod
|
||||
def image_model_combo_option(cls, mode="text"):
|
||||
"""Same as model_combo_option(), for the image-generation node's
|
||||
provider/model DynamicCombo ("text" = text-to-image, "image" = image-to-image)."""
|
||||
if mode == "text":
|
||||
model_ids = cls.list_image_models()
|
||||
build_inputs = cls.image_settings_inputs
|
||||
else:
|
||||
model_ids = cls.image_to_image_model_ids()
|
||||
build_inputs = cls.image_to_image_settings_inputs
|
||||
return io.DynamicCombo.Option(
|
||||
cls.PROVIDER,
|
||||
[
|
||||
io.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
io.DynamicCombo.Option(model_id, build_inputs(model_id))
|
||||
for model_id in model_ids
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def image_bytes_list_from_response(image_response) -> list[bytes]:
|
||||
"""Extract raw image bytes from a litellm ImageResponse: base64-decode
|
||||
each item's b64_json, or fetch its url if that's all the provider gave
|
||||
(e.g. RunwayML's image endpoint only returns a hosted url, never b64)."""
|
||||
import base64
|
||||
|
||||
import httpx
|
||||
|
||||
results = []
|
||||
for item in getattr(image_response, "data", None) or []:
|
||||
b64_json = getattr(item, "b64_json", None)
|
||||
if b64_json:
|
||||
results.append(base64.b64decode(b64_json))
|
||||
continue
|
||||
url = getattr(item, "url", None)
|
||||
if url:
|
||||
response = httpx.get(url)
|
||||
response.raise_for_status()
|
||||
results.append(response.content)
|
||||
continue
|
||||
raise RuntimeError("Image response contained neither b64_json nor url data.")
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def image_bytes_list_to_tensor(image_bytes_list):
|
||||
"""Decode a list of image byte strings (PNG/JPEG/etc) into one batched
|
||||
ComfyUI IMAGE tensor ([B, H, W, C] float32 in [0, 1]).
|
||||
Imports are deferred: torch/numpy/PIL are only available inside a real
|
||||
ComfyUI install, not in this package's own test environment."""
|
||||
from io import BytesIO
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
arrays = []
|
||||
for image_bytes in image_bytes_list:
|
||||
image = Image.open(BytesIO(image_bytes)).convert("RGB")
|
||||
arrays.append(np.array(image).astype(np.float32) / 255.0)
|
||||
return torch.from_numpy(np.stack(arrays, axis=0))
|
||||
+150
-20
@@ -32,7 +32,34 @@ OMNI_STATIC_MODELS = [
|
||||
"gemini-omni-flash-preview",
|
||||
]
|
||||
|
||||
VEO_RESOLUTION_MAP = {Resolutions._1K: "720p", Resolutions._2K: "1080p"}
|
||||
# Imagen: text-to-image only, via litellm's aimage_generation (":predict").
|
||||
IMAGEN_STATIC_MODELS = [
|
||||
"imagen-4.0-generate-001",
|
||||
"imagen-4.0-fast-generate-001",
|
||||
"imagen-3.0-generate-002",
|
||||
]
|
||||
|
||||
# "Nano Banana" family: Gemini's own multimodal image models. Support both
|
||||
# text-to-image (aimage_generation, ":generateContent") and image-to-image
|
||||
# (aimage_edit, also ":generateContent" — Imagen has no edit endpoint at all).
|
||||
# Google's naming: gemini-2.5-flash-image = Nano Banana, gemini-3-pro-image-preview
|
||||
# = Nano Banana Pro, gemini-3.1-flash-image-preview = Nano Banana 2.
|
||||
NANO_BANANA_STATIC_MODELS = [
|
||||
"gemini-2.5-flash-image",
|
||||
"gemini-3-pro-image-preview",
|
||||
"gemini-3.1-flash-image-preview",
|
||||
]
|
||||
|
||||
IMAGE_SIZE_OPTIONS = ["1024x1024", "1792x1024", "1024x1792"]
|
||||
|
||||
|
||||
def _is_nano_banana(model_id: str) -> bool:
|
||||
# Matches litellm's own is_gemini_image_model() dispatch rule exactly —
|
||||
# any image model id containing "gemini" (vs. "imagen") goes through the
|
||||
# generateContent-based, edit-capable code path.
|
||||
return "gemini" in model_id.lower()
|
||||
|
||||
VEO_RESOLUTION_MAP = {Resolutions._720: "720p", Resolutions._1080: "1080p", Resolutions._4K: "4k"}
|
||||
OMNI_RESOLUTION_LABELS = {
|
||||
Resolutions._360: "360p",
|
||||
Resolutions._720: "720p",
|
||||
@@ -40,12 +67,10 @@ OMNI_RESOLUTION_LABELS = {
|
||||
Resolutions._4K: "4k",
|
||||
}
|
||||
|
||||
VEO_VIDEO_SETTINGS = VideoSettings(
|
||||
SIZE=[Sizes._16_9, Sizes._9_16],
|
||||
RESOLUTIONS=[Resolutions._1K, Resolutions._2K],
|
||||
MIN_SECONDS=4,
|
||||
MAX_SECONDS=8,
|
||||
)
|
||||
VEO_SIZE = [Sizes._16_9, Sizes._9_16]
|
||||
VEO_MIN_SECONDS = 4
|
||||
VEO_MAX_SECONDS = 8
|
||||
|
||||
OMNI_VIDEO_SETTINGS = VideoSettings(
|
||||
SIZE=[Sizes._16_9, Sizes._9_16],
|
||||
RESOLUTIONS=[Resolutions._360, Resolutions._720, Resolutions._1080, Resolutions._4K],
|
||||
@@ -58,6 +83,31 @@ def _is_omni(model_id: str) -> bool:
|
||||
return "omni" in model_id.lower()
|
||||
|
||||
|
||||
def _veo_resolutions(model_id: str) -> list[Resolutions]:
|
||||
"""Resolution tiers actually supported per Veo model (verified against
|
||||
Google's Veo docs — not uniform across versions):
|
||||
- Veo 2.0: 720p only.
|
||||
- Veo 3.1 Lite: 720p/1080p, no 4k.
|
||||
- Veo 3.1 (and 3.1 Fast): 720p/1080p/4k.
|
||||
1080p and 4k are additionally locked to 8s-only clips on Google's side;
|
||||
that duration/resolution coupling isn't enforced by this dropdown."""
|
||||
lower = model_id.lower()
|
||||
if "veo-2" in lower:
|
||||
return [Resolutions._720]
|
||||
if "3.1" in lower and "lite" not in lower:
|
||||
return [Resolutions._720, Resolutions._1080, Resolutions._4K]
|
||||
return [Resolutions._720, Resolutions._1080]
|
||||
|
||||
|
||||
def _veo_video_settings(model_id: str) -> VideoSettings:
|
||||
return VideoSettings(
|
||||
SIZE=VEO_SIZE,
|
||||
RESOLUTIONS=_veo_resolutions(model_id),
|
||||
MIN_SECONDS=VEO_MIN_SECONDS,
|
||||
MAX_SECONDS=VEO_MAX_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
class GeminiModel(Model):
|
||||
"""Google Gemini video generation via litellm's 'gemini/<model_id>' provider:
|
||||
Veo (text/image-to-video) and Omni Flash (text-to-video only, so far)."""
|
||||
@@ -103,12 +153,6 @@ class GeminiModel(Model):
|
||||
|
||||
return cls._models_cache
|
||||
|
||||
@classmethod
|
||||
def image_to_video_model_ids(cls) -> list[str]:
|
||||
# Image-to-video is Veo-only for now; Omni Flash's image/video
|
||||
# conditioning isn't implemented yet (see README roadmap).
|
||||
return [model_id for model_id in cls.list_models() if not _is_omni(model_id)]
|
||||
|
||||
@staticmethod
|
||||
def _negative_prompt_input():
|
||||
return io.String.Input(
|
||||
@@ -127,10 +171,11 @@ class GeminiModel(Model):
|
||||
OMNI_VIDEO_SETTINGS.resolution_input(),
|
||||
OMNI_VIDEO_SETTINGS.seconds_input(),
|
||||
]
|
||||
veo_settings = _veo_video_settings(model_id or "")
|
||||
return [
|
||||
VEO_VIDEO_SETTINGS.size_input(),
|
||||
VEO_VIDEO_SETTINGS.resolution_input(),
|
||||
VEO_VIDEO_SETTINGS.seconds_input(),
|
||||
veo_settings.size_input(),
|
||||
veo_settings.resolution_input(),
|
||||
veo_settings.seconds_input(),
|
||||
io.Combo.Input(
|
||||
"person_generation",
|
||||
options=["allow_all", "allow_adult"],
|
||||
@@ -142,10 +187,18 @@ class GeminiModel(Model):
|
||||
|
||||
@classmethod
|
||||
def image_to_video_settings_inputs(cls, model_id: str | None = None):
|
||||
if _is_omni(model_id or ""):
|
||||
return [
|
||||
OMNI_VIDEO_SETTINGS.size_input(),
|
||||
OMNI_VIDEO_SETTINGS.resolution_input(),
|
||||
OMNI_VIDEO_SETTINGS.seconds_input(),
|
||||
reference_image_input(),
|
||||
]
|
||||
veo_settings = _veo_video_settings(model_id or "")
|
||||
return [
|
||||
VEO_VIDEO_SETTINGS.size_input(),
|
||||
VEO_VIDEO_SETTINGS.resolution_input(),
|
||||
VEO_VIDEO_SETTINGS.seconds_input(),
|
||||
veo_settings.size_input(),
|
||||
veo_settings.resolution_input(),
|
||||
veo_settings.seconds_input(),
|
||||
reference_image_input(),
|
||||
io.Combo.Input(
|
||||
"person_generation",
|
||||
@@ -182,7 +235,30 @@ class GeminiModel(Model):
|
||||
if reference_image is None:
|
||||
raise ValueError("Image-to-video requires a reference_image input.")
|
||||
kwargs = cls.video_kwargs(model_id, settings)
|
||||
kwargs["input_reference"] = cls.image_tensor_to_bytesio(reference_image)
|
||||
if _is_omni(model_id):
|
||||
# The Interactions API has no separate "reference image" field like
|
||||
# avideo_generation's input_reference — each image is just another
|
||||
# item in the same ordered `input` content list as the prompt
|
||||
# text, so a batch of N images (built upstream with a Batch
|
||||
# Images node) becomes N image content items, in order.
|
||||
image_contents = [
|
||||
{
|
||||
"type": "image",
|
||||
"data": base64.b64encode(buffer.getvalue()).decode(),
|
||||
"mime_type": "image/png",
|
||||
}
|
||||
for buffer in cls.image_tensor_to_bytesio_list(reference_image)
|
||||
]
|
||||
kwargs["input"] = [{"type": "text", "text": settings.get("prompt", "")}, *image_contents]
|
||||
else:
|
||||
batch_size = getattr(reference_image, "shape", [1])[0]
|
||||
if batch_size > 1:
|
||||
raise ValueError(
|
||||
f"Veo only accepts a single reference image via litellm (got a batch of "
|
||||
f"{batch_size}) — litellm's Veo integration has no multi-image field yet. "
|
||||
"Gemini Omni Flash models accept a batch of reference images if you need that."
|
||||
)
|
||||
kwargs["input_reference"] = cls.image_tensor_to_bytesio(reference_image)
|
||||
return kwargs
|
||||
|
||||
@classmethod
|
||||
@@ -201,3 +277,57 @@ class GeminiModel(Model):
|
||||
"base64 data; downloading a uri-delivered video isn't supported yet."
|
||||
)
|
||||
raise RuntimeError("Gemini Omni response contained no video output.")
|
||||
|
||||
@classmethod
|
||||
def list_image_models(cls) -> list[str]:
|
||||
# Static only, unlike list_models(): there's no safe live-capability
|
||||
# signal to filter on here the way Veo's supported_actions works —
|
||||
# "gemini" appears in the name of ordinary chat models too, so a
|
||||
# live models.list() name-substring filter would wrongly include them.
|
||||
return list(IMAGEN_STATIC_MODELS) + list(NANO_BANANA_STATIC_MODELS)
|
||||
|
||||
@classmethod
|
||||
def image_to_image_model_ids(cls) -> list[str]:
|
||||
# Image-to-image (aimage_edit) only works for the Nano Banana family —
|
||||
# Imagen has no edit endpoint in litellm at all.
|
||||
return [model_id for model_id in cls.list_image_models() if _is_nano_banana(model_id)]
|
||||
|
||||
@classmethod
|
||||
def image_settings_inputs(cls, model_id: str | None = None):
|
||||
return [
|
||||
io.Combo.Input(
|
||||
"size",
|
||||
options=IMAGE_SIZE_OPTIONS,
|
||||
default=IMAGE_SIZE_OPTIONS[0],
|
||||
tooltip="Image dimensions (width x height).",
|
||||
),
|
||||
io.Int.Input(
|
||||
"num_images",
|
||||
default=1,
|
||||
min=1,
|
||||
max=4,
|
||||
tooltip="Number of images to generate at once.",
|
||||
),
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def image_to_image_settings_inputs(cls, model_id: str | None = None):
|
||||
return [
|
||||
*cls.image_settings_inputs(model_id),
|
||||
reference_image_input(),
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def image_kwargs(cls, model_id: str, settings: dict) -> dict:
|
||||
return {"size": settings["size"], "n": settings["num_images"]}
|
||||
|
||||
@classmethod
|
||||
def image_to_image_kwargs(cls, model_id: str, settings: dict) -> dict:
|
||||
reference_image = settings.get("reference_image")
|
||||
if reference_image is None:
|
||||
raise ValueError("Image-to-image requires a reference_image input.")
|
||||
kwargs = cls.image_kwargs(model_id, settings)
|
||||
# litellm's aimage_edit takes a single image or a list — a batch of N
|
||||
# (built upstream with a Batch Images node) becomes a list of N.
|
||||
kwargs["image"] = cls.image_tensor_to_bytesio_list(reference_image)
|
||||
return kwargs
|
||||
|
||||
@@ -17,7 +17,10 @@ class RunwayMLModel(Model):
|
||||
|
||||
@classmethod
|
||||
def list_models(cls) -> list[str]:
|
||||
return ["gen3a_turbo", "gen4_turbo", "gen4_aleph", "gen4_image", "gen4_image_turbo"]
|
||||
# gen4_image/gen4_image_turbo are image models, not video — they don't
|
||||
# belong in this list (see README roadmap: not yet wired up as image
|
||||
# models on the image-generation node either).
|
||||
return ["gen3a_turbo", "gen4_turbo", "gen4_aleph"]
|
||||
|
||||
@classmethod
|
||||
def video_settings_inputs(cls, model_id: str | None = None):
|
||||
|
||||
+112
-3
@@ -2,10 +2,11 @@ import asyncio
|
||||
from io import BytesIO
|
||||
|
||||
from comfy_api.latest import ComfyAPI, InputImpl, io
|
||||
from litellm import avideo_content, avideo_generation, avideo_status
|
||||
from litellm import aimage_edit, aimage_generation, avideo_content, avideo_generation, avideo_status
|
||||
from litellm.interactions import acreate as acreate_interaction
|
||||
|
||||
from .models.azure import AzureModel
|
||||
from .models.base import Model
|
||||
from .models.gemini import GeminiModel
|
||||
from .models.openai import OpenAIModel
|
||||
from .models.runwayml import RunwayMLModel
|
||||
@@ -22,6 +23,12 @@ _PROVIDERS_BY_NAME = {p.PROVIDER: p for p in _PROVIDERS}
|
||||
_IMAGE_TO_VIDEO_PROVIDERS = [GeminiModel]
|
||||
_IMAGE_TO_VIDEO_PROVIDERS_BY_NAME = {p.PROVIDER: p for p in _IMAGE_TO_VIDEO_PROVIDERS}
|
||||
|
||||
# Image generation (separate node, below): Gemini only for now (Imagen +
|
||||
# Nano Banana). Other providers can be added once they implement
|
||||
# list_image_models()/image_settings_inputs()/image_kwargs().
|
||||
_IMAGE_PROVIDERS = [GeminiModel]
|
||||
_IMAGE_PROVIDERS_BY_NAME = {p.PROVIDER: p for p in _IMAGE_PROVIDERS}
|
||||
|
||||
|
||||
class APIVideoGenerate(io.ComfyNode):
|
||||
"""Text-to-video and image-to-video generation via litellm, using the caller's own provider API key (BYOK)."""
|
||||
@@ -118,6 +125,7 @@ class APIVideoGenerate(io.ComfyNode):
|
||||
model_dict = provider_dict["model"]
|
||||
model_id = model_dict["model"]
|
||||
settings = {k: v for k, v in model_dict.items() if k != "model"}
|
||||
settings["prompt"] = prompt
|
||||
|
||||
if source_mode == "text":
|
||||
provider_cls = _PROVIDERS_BY_NAME[provider_name]
|
||||
@@ -133,8 +141,11 @@ class APIVideoGenerate(io.ComfyNode):
|
||||
# report. Tick the bar up by elapsed/timeout while we wait on it,
|
||||
# the same heuristic the polling path below uses between polls,
|
||||
# instead of sitting at 0% and then jumping straight to done.
|
||||
interaction_input = extra_kwargs.pop("input", prompt)
|
||||
request = asyncio.ensure_future(
|
||||
acreate_interaction(model=model, input=prompt, timeout=timeout, **auth_kwargs, **extra_kwargs)
|
||||
acreate_interaction(
|
||||
model=model, input=interaction_input, timeout=timeout, **auth_kwargs, **extra_kwargs
|
||||
)
|
||||
)
|
||||
elapsed = 0.0
|
||||
await api.execution.set_progress(0.0, timeout, node_id=node_id)
|
||||
@@ -170,4 +181,102 @@ class APIVideoGenerate(io.ComfyNode):
|
||||
return io.NodeOutput(InputImpl.VideoFromFile(BytesIO(video_bytes)))
|
||||
|
||||
|
||||
NODES = [APIVideoGenerate]
|
||||
class APIImageGenerate(io.ComfyNode):
|
||||
"""Text-to-image and image-to-image generation via litellm, using the caller's own provider API key (BYOK)."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="APIImageGenerate",
|
||||
display_name="API Image Generate (BYOK)",
|
||||
category="external_api/image",
|
||||
description=(
|
||||
"Generates an image from a text prompt (or a text prompt plus one or more "
|
||||
"reference images) via litellm, using your own provider API key. "
|
||||
"Pick a source, provider, and model to reveal their parameters."
|
||||
),
|
||||
inputs=[
|
||||
io.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
tooltip="Text description of the image to generate.",
|
||||
),
|
||||
io.DynamicCombo.Input(
|
||||
"source",
|
||||
options=[
|
||||
io.DynamicCombo.Option(
|
||||
"text",
|
||||
[
|
||||
io.DynamicCombo.Input(
|
||||
"provider",
|
||||
options=[
|
||||
p.image_model_combo_option(mode="text") for p in _IMAGE_PROVIDERS
|
||||
],
|
||||
tooltip="Image generation provider.",
|
||||
),
|
||||
],
|
||||
),
|
||||
io.DynamicCombo.Option(
|
||||
"image",
|
||||
[
|
||||
io.DynamicCombo.Input(
|
||||
"provider",
|
||||
options=[
|
||||
p.image_model_combo_option(mode="image") for p in _IMAGE_PROVIDERS
|
||||
],
|
||||
tooltip="Image-to-image provider.",
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Generation source: a text prompt alone, or a text prompt plus reference image(s).",
|
||||
),
|
||||
io.String.Input(
|
||||
"api_key",
|
||||
default="",
|
||||
optional=True,
|
||||
tooltip=(
|
||||
"Overrides the provider's standard env var (e.g. GEMINI_API_KEY) when set. "
|
||||
"Warning: unlike an env var, this value is saved into the workflow JSON and "
|
||||
"into output file metadata. Prefer leaving this blank and setting the env "
|
||||
"var on the ComfyUI server instead."
|
||||
),
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(),
|
||||
],
|
||||
hidden=[io.Hidden.unique_id],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(cls, prompt, source, api_key=""):
|
||||
auth_kwargs = {"api_key": api_key} if api_key else {}
|
||||
|
||||
source_mode = source["source"]
|
||||
provider_dict = source["provider"]
|
||||
provider_name = provider_dict["provider"]
|
||||
model_dict = provider_dict["model"]
|
||||
model_id = model_dict["model"]
|
||||
settings = {k: v for k, v in model_dict.items() if k != "model"}
|
||||
settings["prompt"] = prompt
|
||||
|
||||
provider_cls = _IMAGE_PROVIDERS_BY_NAME[provider_name]
|
||||
model = f"{provider_name}/{model_id}"
|
||||
|
||||
if source_mode == "text":
|
||||
extra_kwargs = provider_cls.image_kwargs(model_id, settings)
|
||||
response = await aimage_generation(model=model, prompt=prompt, **auth_kwargs, **extra_kwargs)
|
||||
else:
|
||||
extra_kwargs = provider_cls.image_to_image_kwargs(model_id, settings)
|
||||
image = extra_kwargs.pop("image")
|
||||
response = await aimage_edit(
|
||||
model=model, prompt=prompt, image=image, **auth_kwargs, **extra_kwargs
|
||||
)
|
||||
|
||||
image_bytes_list = Model.image_bytes_list_from_response(response)
|
||||
tensor = Model.image_bytes_list_to_tensor(image_bytes_list)
|
||||
return io.NodeOutput(tensor)
|
||||
|
||||
|
||||
NODES = [APIVideoGenerate, APIImageGenerate]
|
||||
|
||||
@@ -33,7 +33,12 @@ def reference_image_input(id="reference_image", optional=True):
|
||||
return io.Image.Input(
|
||||
id,
|
||||
optional=optional,
|
||||
tooltip="Reference image to seed the video. Connect a Load Image node or any other node's IMAGE output.",
|
||||
tooltip=(
|
||||
"Reference image(s) to seed the video. Connect a Load Image node or any other "
|
||||
"node's IMAGE output. Gemini Omni Flash accepts a batch (e.g. via a Batch Images "
|
||||
"node) as multiple reference images, in order; other models (e.g. Veo) accept only "
|
||||
"a single image and will raise a clear error if given a batch of more than one."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
||||
+1
-1
@@ -45,7 +45,7 @@ def _install_fake_comfy_api():
|
||||
class Combo(_InputFactory):
|
||||
pass
|
||||
|
||||
class Image(_InputFactory):
|
||||
class Image(_InputFactory, _OutputFactory):
|
||||
pass
|
||||
|
||||
class _DynamicComboOption:
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
from src.models.gemini import GeminiModel
|
||||
|
||||
|
||||
def _resolution_options(model_id):
|
||||
inputs = GeminiModel.video_settings_inputs(model_id)
|
||||
resolution_input = next(i for i in inputs if i.name == "resolution")
|
||||
return resolution_input.kwargs["options"]
|
||||
|
||||
|
||||
def test_veo_2_0_only_supports_720p():
|
||||
assert _resolution_options("veo-2.0-generate-001") == ["720"]
|
||||
|
||||
|
||||
def test_veo_3_1_lite_has_no_4k():
|
||||
assert _resolution_options("veo-3.1-lite-generate-preview") == ["720", "1080"]
|
||||
|
||||
|
||||
def test_veo_3_1_supports_4k():
|
||||
assert _resolution_options("veo-3.1-generate-preview") == ["720", "1080", "4K"]
|
||||
assert _resolution_options("veo-3.1-fast-generate-preview") == ["720", "1080", "4K"]
|
||||
|
||||
|
||||
def test_veo_resolution_kwargs_maps_4k():
|
||||
kwargs = GeminiModel.video_kwargs(
|
||||
"veo-3.1-generate-preview",
|
||||
{"size": "16:9", "resolution": "4K", "seconds": 8, "person_generation": "allow_all"},
|
||||
)
|
||||
assert kwargs["resolution"] == "4k"
|
||||
@@ -0,0 +1,127 @@
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from src import nodes
|
||||
|
||||
|
||||
class FakeImageObject:
|
||||
def __init__(self, b64_json=None, url=None):
|
||||
self.b64_json = b64_json
|
||||
self.url = url
|
||||
|
||||
|
||||
class FakeImageResponse:
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
|
||||
|
||||
def _text_source(model="imagen-4.0-generate-001", **overrides):
|
||||
settings = {"model": model, "size": "1024x1024", "num_images": 1}
|
||||
settings.update(overrides)
|
||||
return {"source": "text", "provider": {"provider": "gemini", "model": settings}}
|
||||
|
||||
|
||||
def _image_source(model="gemini-2.5-flash-image", **overrides):
|
||||
settings = {"model": model, "size": "1024x1024", "num_images": 1}
|
||||
settings.update(overrides)
|
||||
return {"source": "image", "provider": {"provider": "gemini", "model": settings}}
|
||||
|
||||
|
||||
def test_schema_lists_only_nano_banana_for_image_source():
|
||||
schema = nodes.APIImageGenerate.define_schema()
|
||||
source_input = next(i for i in schema.inputs if i.name == "source")
|
||||
text_option, image_option = source_input.options
|
||||
|
||||
text_provider_input = text_option.inputs[0]
|
||||
gemini_text_option = next(o for o in text_provider_input.options if o.key == "gemini")
|
||||
text_model_ids = {o.key for o in gemini_text_option.inputs[0].options}
|
||||
assert "imagen-4.0-generate-001" in text_model_ids
|
||||
assert "gemini-2.5-flash-image" in text_model_ids
|
||||
|
||||
image_provider_input = image_option.inputs[0]
|
||||
gemini_image_option = next(o for o in image_provider_input.options if o.key == "gemini")
|
||||
image_model_ids = {o.key for o in gemini_image_option.inputs[0].options}
|
||||
assert "gemini-2.5-flash-image" in image_model_ids
|
||||
assert "imagen-4.0-generate-001" not in image_model_ids
|
||||
|
||||
for model_option in gemini_image_option.inputs[0].options:
|
||||
field_names = {i.name for i in model_option.inputs}
|
||||
assert {"size", "num_images", "reference_image"} <= field_names
|
||||
|
||||
|
||||
def test_text_to_image_calls_aimage_generation():
|
||||
async def run():
|
||||
gen_mock = AsyncMock(
|
||||
return_value=FakeImageResponse(data=[FakeImageObject(b64_json="ZmFrZS1wbmc=")])
|
||||
)
|
||||
|
||||
with patch.object(nodes, "aimage_generation", gen_mock), patch.object(
|
||||
nodes.Model,
|
||||
"image_bytes_list_to_tensor",
|
||||
return_value="fake-tensor",
|
||||
) as tensor_mock:
|
||||
result = await nodes.APIImageGenerate.execute(
|
||||
prompt="a cat wearing a hat",
|
||||
source=_text_source(),
|
||||
)
|
||||
|
||||
gen_mock.assert_awaited_once()
|
||||
assert gen_mock.call_args.kwargs["model"] == "gemini/imagen-4.0-generate-001"
|
||||
assert gen_mock.call_args.kwargs["prompt"] == "a cat wearing a hat"
|
||||
assert gen_mock.call_args.kwargs["size"] == "1024x1024"
|
||||
assert gen_mock.call_args.kwargs["n"] == 1
|
||||
tensor_mock.assert_called_once_with([b"fake-png"])
|
||||
assert result.args[0] == "fake-tensor"
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_image_to_image_calls_aimage_edit_with_batched_images():
|
||||
async def run():
|
||||
edit_mock = AsyncMock(
|
||||
return_value=FakeImageResponse(data=[FakeImageObject(b64_json="ZmFrZS1wbmc=")])
|
||||
)
|
||||
fake_tensor = object()
|
||||
fake_buffers = [object(), object()]
|
||||
|
||||
source = _image_source()
|
||||
source["provider"]["model"]["reference_image"] = fake_tensor
|
||||
|
||||
with patch.object(nodes, "aimage_edit", edit_mock), patch.object(
|
||||
nodes.GeminiModel,
|
||||
"image_tensor_to_bytesio_list",
|
||||
return_value=fake_buffers,
|
||||
) as split_mock, patch.object(
|
||||
nodes.Model, "image_bytes_list_to_tensor", return_value="fake-tensor"
|
||||
):
|
||||
result = await nodes.APIImageGenerate.execute(
|
||||
prompt="combine these two",
|
||||
source=source,
|
||||
)
|
||||
|
||||
split_mock.assert_called_once_with(fake_tensor)
|
||||
edit_mock.assert_awaited_once()
|
||||
assert edit_mock.call_args.kwargs["model"] == "gemini/gemini-2.5-flash-image"
|
||||
assert edit_mock.call_args.kwargs["image"] == fake_buffers
|
||||
assert edit_mock.call_args.kwargs["prompt"] == "combine these two"
|
||||
assert result.args[0] == "fake-tensor"
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_image_to_image_requires_reference_image():
|
||||
async def run():
|
||||
edit_mock = AsyncMock()
|
||||
|
||||
with patch.object(nodes, "aimage_edit", edit_mock):
|
||||
with pytest.raises(ValueError, match="reference_image"):
|
||||
await nodes.APIImageGenerate.execute(
|
||||
prompt="anything",
|
||||
source=_image_source(),
|
||||
)
|
||||
|
||||
edit_mock.assert_not_awaited()
|
||||
|
||||
asyncio.run(run())
|
||||
+164
-11
@@ -1,6 +1,7 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import types
|
||||
from io import BytesIO
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
@@ -19,7 +20,7 @@ def _gemini_source(model="veo-3.1-fast-generate-preview", **overrides):
|
||||
settings = {
|
||||
"model": model,
|
||||
"size": "16:9",
|
||||
"resolution": "1K",
|
||||
"resolution": "720",
|
||||
"seconds": 8,
|
||||
"person_generation": "allow_all",
|
||||
}
|
||||
@@ -219,15 +220,25 @@ def test_schema_builds_nested_source_provider_model_combo():
|
||||
gemini_model_input = image_provider_input.options[0].inputs[0]
|
||||
assert gemini_model_input.name == "model"
|
||||
gemini_image_model_ids = {o.key for o in gemini_model_input.options}
|
||||
assert "gemini-omni-flash-preview" not in gemini_image_model_ids
|
||||
for model_option in gemini_model_input.options:
|
||||
field_names = {i.name for i in model_option.inputs}
|
||||
assert {"size", "resolution", "seconds", "reference_image", "person_generation"} <= field_names
|
||||
assert "veo-3.1-fast-generate-preview" in gemini_image_model_ids
|
||||
assert "gemini-omni-flash-preview" in gemini_image_model_ids
|
||||
|
||||
person_generation_input = next(
|
||||
i for i in model_option.inputs if i.name == "person_generation"
|
||||
)
|
||||
assert person_generation_input.kwargs["options"] == ["allow_adult"]
|
||||
veo_image_option = next(
|
||||
o for o in gemini_model_input.options if o.key == "veo-3.1-fast-generate-preview"
|
||||
)
|
||||
veo_image_field_names = {i.name for i in veo_image_option.inputs}
|
||||
assert {"size", "resolution", "seconds", "reference_image", "person_generation"} <= veo_image_field_names
|
||||
person_generation_input = next(
|
||||
i for i in veo_image_option.inputs if i.name == "person_generation"
|
||||
)
|
||||
assert person_generation_input.kwargs["options"] == ["allow_adult"]
|
||||
|
||||
omni_image_option = next(
|
||||
o for o in gemini_model_input.options if o.key == "gemini-omni-flash-preview"
|
||||
)
|
||||
omni_image_field_names = {i.name for i in omni_image_option.inputs}
|
||||
assert {"size", "resolution", "seconds", "reference_image"} <= omni_image_field_names
|
||||
assert "person_generation" not in omni_image_field_names
|
||||
|
||||
|
||||
def test_gemini_image_to_video_kwargs_mapping():
|
||||
@@ -245,7 +256,7 @@ def test_gemini_image_to_video_kwargs_mapping():
|
||||
"model": {
|
||||
"model": "veo-3.1-fast-generate-preview",
|
||||
"size": "9:16",
|
||||
"resolution": "2K",
|
||||
"resolution": "1080",
|
||||
"seconds": 8,
|
||||
"reference_image": fake_tensor,
|
||||
"person_generation": "allow_adult",
|
||||
@@ -283,7 +294,7 @@ def test_gemini_image_to_video_requires_reference_image():
|
||||
"model": {
|
||||
"model": "veo-3.1-fast-generate-preview",
|
||||
"size": "9:16",
|
||||
"resolution": "2K",
|
||||
"resolution": "1080",
|
||||
"seconds": 8,
|
||||
"person_generation": "allow_adult",
|
||||
},
|
||||
@@ -381,6 +392,148 @@ def test_gemini_omni_uri_delivery_raises_clear_error():
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_gemini_omni_image_to_video_sends_text_and_image_content():
|
||||
async def run():
|
||||
interaction_mock = AsyncMock(
|
||||
return_value=FakeInteractionResponse(
|
||||
steps=[
|
||||
{
|
||||
"type": "model_output",
|
||||
"content": [
|
||||
{
|
||||
"type": "video",
|
||||
"data": base64.b64encode(b"fake-mp4-bytes").decode(),
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
)
|
||||
)
|
||||
fake_tensor = object()
|
||||
fake_bytesio_content = b"fake-png-bytes"
|
||||
|
||||
source = {
|
||||
"source": "image",
|
||||
"provider": {
|
||||
"provider": "gemini",
|
||||
"model": {
|
||||
"model": "gemini-omni-flash-preview",
|
||||
"size": "16:9",
|
||||
"resolution": "720",
|
||||
"seconds": 5,
|
||||
"reference_image": fake_tensor,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
with patch.object(nodes, "acreate_interaction", interaction_mock), patch.object(
|
||||
nodes.GeminiModel,
|
||||
"image_tensor_to_bytesio_list",
|
||||
return_value=[BytesIO(fake_bytesio_content)],
|
||||
) as tensor_mock:
|
||||
result = await nodes.APIVideoGenerate.execute(
|
||||
prompt="make the person in the photo wave",
|
||||
source=source,
|
||||
)
|
||||
|
||||
tensor_mock.assert_called_once_with(fake_tensor)
|
||||
kwargs = interaction_mock.call_args.kwargs
|
||||
assert kwargs["model"] == "gemini/gemini-omni-flash-preview"
|
||||
assert kwargs["input"] == [
|
||||
{"type": "text", "text": "make the person in the photo wave"},
|
||||
{
|
||||
"type": "image",
|
||||
"data": base64.b64encode(fake_bytesio_content).decode(),
|
||||
"mime_type": "image/png",
|
||||
},
|
||||
]
|
||||
assert kwargs["response_format"] == {
|
||||
"type": "video",
|
||||
"aspect_ratio": "16:9",
|
||||
"resolution": "720p",
|
||||
}
|
||||
|
||||
video = result.args[0]
|
||||
assert video.file_obj.getvalue() == b"fake-mp4-bytes"
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_gemini_omni_image_to_video_sends_one_content_item_per_batched_image():
|
||||
async def run():
|
||||
interaction_mock = AsyncMock(
|
||||
return_value=FakeInteractionResponse(
|
||||
steps=[
|
||||
{
|
||||
"type": "model_output",
|
||||
"content": [{"type": "video", "data": base64.b64encode(b"vid").decode()}],
|
||||
}
|
||||
]
|
||||
)
|
||||
)
|
||||
fake_tensor = object()
|
||||
frame_bytes = [b"frame-1", b"frame-2", b"frame-3"]
|
||||
|
||||
source = {
|
||||
"source": "image",
|
||||
"provider": {
|
||||
"provider": "gemini",
|
||||
"model": {
|
||||
"model": "gemini-omni-flash-preview",
|
||||
"size": "16:9",
|
||||
"resolution": "720",
|
||||
"seconds": 5,
|
||||
"reference_image": fake_tensor,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
with patch.object(nodes, "acreate_interaction", interaction_mock), patch.object(
|
||||
nodes.GeminiModel,
|
||||
"image_tensor_to_bytesio_list",
|
||||
return_value=[BytesIO(b) for b in frame_bytes],
|
||||
):
|
||||
await nodes.APIVideoGenerate.execute(prompt="combine these", source=source)
|
||||
|
||||
kwargs = interaction_mock.call_args.kwargs
|
||||
image_items = [item for item in kwargs["input"] if item["type"] == "image"]
|
||||
assert len(image_items) == 3
|
||||
assert [item["data"] for item in image_items] == [
|
||||
base64.b64encode(b).decode() for b in frame_bytes
|
||||
]
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_veo_image_to_video_rejects_batched_reference_image():
|
||||
async def run():
|
||||
gen_mock = AsyncMock()
|
||||
|
||||
fake_tensor = types.SimpleNamespace(shape=(3, 64, 64, 3))
|
||||
source = {
|
||||
"source": "image",
|
||||
"provider": {
|
||||
"provider": "gemini",
|
||||
"model": {
|
||||
"model": "veo-3.1-fast-generate-preview",
|
||||
"size": "16:9",
|
||||
"resolution": "720",
|
||||
"seconds": 8,
|
||||
"reference_image": fake_tensor,
|
||||
"person_generation": "allow_adult",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
with patch.object(nodes, "avideo_generation", gen_mock):
|
||||
with pytest.raises(ValueError, match="single reference image"):
|
||||
await nodes.APIVideoGenerate.execute(prompt="anything", source=source)
|
||||
|
||||
gen_mock.assert_not_awaited()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_azure_provider_string_uses_azure_prefix():
|
||||
async def run():
|
||||
gen_mock = AsyncMock(return_value=FakeVideoObject("completed"))
|
||||
|
||||
Reference in New Issue
Block a user