diff --git a/README.md b/README.md index 1e38e5c..58242bc 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/src/models/base.py b/src/models/base.py index 3970b4a..a673ee7 100644 --- a/src/models/base.py +++ b/src/models/base.py @@ -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.") \ No newline at end of file + @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)) \ No newline at end of file diff --git a/src/models/gemini.py b/src/models/gemini.py index bc43878..e8a8ea9 100644 --- a/src/models/gemini.py +++ b/src/models/gemini.py @@ -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/' 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 diff --git a/src/models/runwayml.py b/src/models/runwayml.py index 2d0eae9..6ec67ab 100644 --- a/src/models/runwayml.py +++ b/src/models/runwayml.py @@ -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): diff --git a/src/nodes.py b/src/nodes.py index 4836307..0cd68b3 100644 --- a/src/nodes.py +++ b/src/nodes.py @@ -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] diff --git a/src/settings/video.py b/src/settings/video.py index eb1997a..e5e9f11 100644 --- a/src/settings/video.py +++ b/src/settings/video.py @@ -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." + ), ) diff --git a/tests/conftest.py b/tests/conftest.py index 7fb5f40..e182d22 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -45,7 +45,7 @@ def _install_fake_comfy_api(): class Combo(_InputFactory): pass - class Image(_InputFactory): + class Image(_InputFactory, _OutputFactory): pass class _DynamicComboOption: diff --git a/tests/test_gemini_resolutions.py b/tests/test_gemini_resolutions.py new file mode 100644 index 0000000..0f8cc36 --- /dev/null +++ b/tests/test_gemini_resolutions.py @@ -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" diff --git a/tests/test_image_nodes.py b/tests/test_image_nodes.py new file mode 100644 index 0000000..bf5e894 --- /dev/null +++ b/tests/test_image_nodes.py @@ -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()) diff --git a/tests/test_nodes.py b/tests/test_nodes.py index 2847404..38271c8 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -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"))