From 3ea918e84c0ab687802aeb94e94ca070afb9e64e Mon Sep 17 00:00:00 2001 From: Andy Ratsirarson Date: Mon, 7 Sep 2026 20:44:24 -0700 Subject: [PATCH] Add Gemini Omni Flash support and fix stale Veo model listings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Merge Omni Flash into the existing "gemini" provider: one dropdown entry, model list now includes both Veo and Omni Flash ids, each revealing its own correct settings (resolution/duration/etc). - Omni Flash runs through litellm's request/response Interactions API (litellm.interactions.acreate) instead of the polling avideo_generation/avideo_status/avideo_content trio Veo and the other providers use, since it has no status-polling endpoint. Text-to-video only for now; image/video-conditioned input and conversational editing aren't implemented yet. - Fix progress reporting for the Omni path: it now ticks up by elapsed/timeout while waiting on the single request instead of sitting at 0% and jumping straight to done. - Fix live Veo model listing to filter out models Google's API no longer serves for generation (checks supported_actions for predictLongRunning) rather than trusting name-matching alone — root cause of a live 404 on veo-3.1-fast-generate-001, which is also dropped from the static fallback list. - Bump litellm 1.92.0 -> 1.100.0. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01CVFBqa77Zo9t4ps8wt7YB2 --- README.md | 10 ++- pyproject.toml | 2 +- requirements.txt | 2 +- src/models/azure.py | 2 +- src/models/base.py | 38 ++++++-- src/models/gemini.py | 145 +++++++++++++++++++++++++------ src/models/openai.py | 6 +- src/models/runwayml.py | 8 +- src/nodes.py | 29 ++++++- src/settings/base.py | 1 + src/settings/image.py | 2 +- src/settings/video.py | 2 +- tests/test_gemini_list_models.py | 88 +++++++++++++++++++ tests/test_nodes.py | 99 +++++++++++++++++++++ 14 files changed, 383 insertions(+), 51 deletions(-) create mode 100644 tests/test_gemini_list_models.py diff --git a/README.md b/README.md index e379291..1e38e5c 100644 --- a/README.md +++ b/README.md @@ -14,14 +14,16 @@ BYOK (bring-your-own-key) ComfyUI nodes for generating video through external pr | Provider | `provider` value | Models | Aspect ratios | Resolutions | Duration | |---|---|---|---|---|---| -| Google Gemini (Veo) | `gemini` | Fetched live from your account via `google-genai`; falls back to a built-in list (`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`) if that fails | 16:9, 9:16 | 1K, 2K | 4–8s | +| 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 | | 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 | \* OpenAI/Azure's upper duration bound isn't published by the provider; 12s is a conservative default rather than a confirmed hard limit. -Gemini's model list is fetched live from `google.genai.Client().models.list()` at node-registration time when a Gemini API key is available on the ComfyUI server, so new Veo releases show up automatically without an update to this package. +Gemini's model list is fetched live from `google.genai.Client().models.list()` at node-registration time when a Gemini API key is available on the ComfyUI server, so new Veo/Omni releases show up automatically without an update to this package. The single `gemini` provider's `model` dropdown lists both families together — picking a Veo model reveals Veo's settings, picking an Omni Flash model reveals Omni's. + +**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. @@ -29,7 +31,7 @@ Gemini's model list is fetched live from `google.genai.Client().models.list()` a - ComfyUI with the V3 node API (`comfy_api.latest`). - Python dependencies (installed automatically, see `requirements.txt`): - - `litellm==1.92.0` + - `litellm==1.100.0` - `google-genai==1.47.0` - An API key from at least one of the providers above. @@ -94,6 +96,8 @@ This currently supports exactly **one** reference image per generation. It does - **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. +- **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/pyproject.toml b/pyproject.toml index deb4ef8..7146d83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,7 @@ license = "MIT" # "Environment :: GPU :: Apple Metal", # Apple Metal support # ] -dependencies = ["litellm==1.92.0", "google-genai==1.47.0"] +dependencies = ["litellm==1.100.0", "google-genai==1.47.0"] [project.optional-dependencies] dev = ["pytest>=7.0"] diff --git a/requirements.txt b/requirements.txt index c0dec69..a95d6fe 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,2 @@ -litellm==1.92.0 +litellm==1.100.0 google-genai==1.47.0 \ No newline at end of file diff --git a/src/models/azure.py b/src/models/azure.py index c21d2d1..d87204f 100644 --- a/src/models/azure.py +++ b/src/models/azure.py @@ -1,4 +1,4 @@ -from src.models.openai import _SoraModel +from .openai import _SoraModel class AzureModel(_SoraModel): diff --git a/src/models/base.py b/src/models/base.py index 31638a7..3970b4a 100644 --- a/src/models/base.py +++ b/src/models/base.py @@ -12,13 +12,29 @@ class Model: PROVIDER: str = None VIDEO_SETTINGS = None + @classmethod + def uses_interactions_api(cls, model_id: str | None = None) -> bool: + """True if this model must be called through litellm's request/response + Interactions API (litellm.interactions.acreate) instead of the polling + avideo_generation/avideo_status/avideo_content trio. A provider class + can cover models that need either (e.g. Gemini's Veo vs. Omni Flash), + so this is decided per model_id rather than being a fixed class flag.""" + return False + @classmethod def list_models(cls) -> list[str]: """Return this provider's available model ids for text-to-video.""" raise NotImplementedError(f"list_models not implemented for {cls.__name__}.") @classmethod - def video_settings_inputs(cls): + def image_to_video_model_ids(cls) -> list[str]: + """Model ids from list_models() that also support image-to-video. + Defaults to all of them; override when only some do (e.g. Gemini's + Veo models support it, but its Omni Flash models don't yet).""" + return cls.list_models() + + @classmethod + def video_settings_inputs(cls, model_id: str | None = None): """Build the io.Input list revealed once a specific model is selected.""" if cls.VIDEO_SETTINGS is None: raise NotImplementedError(f"{cls.__name__} does not support video generation.") @@ -36,7 +52,7 @@ class Model: raise NotImplementedError(f"video_kwargs not implemented for {cls.__name__}.") @classmethod - def image_to_video_settings_inputs(cls): + def image_to_video_settings_inputs(cls, model_id: str | None = None): """Build the io.Input list revealed once a specific model is selected, for the image-to-video ('image' source) path. Distinct from video_settings_inputs() because providers can genuinely differ here @@ -50,6 +66,13 @@ class Model: expects for image-to-video.""" raise NotImplementedError(f"image_to_video_kwargs not implemented for {cls.__name__}.") + @classmethod + def extract_video_bytes(cls, response) -> bytes: + """USES_INTERACTIONS_API providers only: pull the generated video's raw + bytes out of the litellm InteractionsAPIResponse returned by + litellm.interactions.acreate().""" + raise NotImplementedError(f"extract_video_bytes not implemented for {cls.__name__}.") + @staticmethod def image_tensor_to_bytesio(image_tensor): """Convert a ComfyUI IMAGE tensor ([B, H, W, C] float32 in [0, 1]) into @@ -74,15 +97,20 @@ class Model: """Build this provider's DynamicCombo.Option: a nested 'model' combo where each model reveals this provider's settings for the given source mode ("text" = text-to-video, "image" = image-to-video).""" - build_inputs = cls.video_settings_inputs if mode == "text" else cls.image_to_video_settings_inputs + if mode == "text": + model_ids = cls.list_models() + build_inputs = cls.video_settings_inputs + else: + model_ids = cls.image_to_video_model_ids() + build_inputs = cls.image_to_video_settings_inputs return io.DynamicCombo.Option( cls.PROVIDER, [ io.DynamicCombo.Input( "model", options=[ - io.DynamicCombo.Option(model_id, build_inputs()) - for model_id in cls.list_models() + io.DynamicCombo.Option(model_id, build_inputs(model_id)) + for model_id in model_ids ], ), ], diff --git a/src/models/gemini.py b/src/models/gemini.py index 3c75c49..bc43878 100644 --- a/src/models/gemini.py +++ b/src/models/gemini.py @@ -1,34 +1,75 @@ +import base64 + from comfy_api.latest import io -from src.models.base import Model -from src.settings.base import Resolutions, Sizes -from src.settings.video import VideoSettings, reference_image_input +from .base import Model +from ..settings.base import Resolutions, Sizes +from ..settings.video import VideoSettings, reference_image_input -STATIC_MODELS = [ +# Veo (text/image-to-video, polled via avideo_generation) and Omni Flash +# (text-to-video only so far, via litellm's request/response Interactions +# API) are both under the single "gemini" provider dropdown entry, and +# dispatched on per model_id below — they share Google's account/API key but +# are otherwise unrelated APIs with different parameters and response shapes. +VEO_STATIC_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", ] -RESOLUTION_MAP = {Resolutions._1K: "720p", Resolutions._2K: "1080p"} +# The REST method Veo generation calls (litellm's avideo_generation -> +# POST {model}:predictLongRunning). Google's model list reports, per model, +# which methods it actually supports right now — model names alone aren't +# reliable, since Google retires/moves specific ids off this method without +# renaming them (e.g. veo-3.1-fast-generate-001 404s here as of 2026-09 while +# still resolving as a real model name). +VEO_GENERATION_ACTION = "predictLongRunning" + +OMNI_STATIC_MODELS = [ + "gemini-omni-1.1-flash", + "gemini-omni-flash-preview", +] + +VEO_RESOLUTION_MAP = {Resolutions._1K: "720p", Resolutions._2K: "1080p"} +OMNI_RESOLUTION_LABELS = { + Resolutions._360: "360p", + Resolutions._720: "720p", + Resolutions._1080: "1080p", + 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, +) +OMNI_VIDEO_SETTINGS = VideoSettings( + SIZE=[Sizes._16_9, Sizes._9_16], + RESOLUTIONS=[Resolutions._360, Resolutions._720, Resolutions._1080, Resolutions._4K], + MIN_SECONDS=3, + MAX_SECONDS=10, +) + + +def _is_omni(model_id: str) -> bool: + return "omni" in model_id.lower() class GeminiModel(Model): - """Google Veo text-to-video via litellm's 'gemini/' provider.""" + """Google Gemini video generation via litellm's 'gemini/' provider: + Veo (text/image-to-video) and Omni Flash (text-to-video only, so far).""" PROVIDER = "gemini" - VIDEO_SETTINGS = VideoSettings( - SIZE=[Sizes._16_9, Sizes._9_16], - RESOLUTIONS=[Resolutions._1K, Resolutions._2K], - MIN_SECONDS=4, - MAX_SECONDS=8, - ) _models_cache = None + @classmethod + def uses_interactions_api(cls, model_id: str | None = None) -> bool: + return _is_omni(model_id or "") + @classmethod def list_models(cls) -> list[str]: if cls._models_cache is not None: @@ -43,13 +84,31 @@ class GeminiModel(Model): if name.startswith("models/"): name = name[len("models/"):] if "veo" in name.lower(): + # Only offer it if Google's own listing says it still + # supports video generation. A model can resolve as a + # real name yet 404 on the actual generate call once + # Google retires it from this method (unknown/missing + # supported_actions is treated as "don't know" rather + # than excluded, so a leaner API response can't hide + # every model). + actions = getattr(model, "supported_actions", None) + if actions and VEO_GENERATION_ACTION not in actions: + continue models.append(name) - cls._models_cache = models or list(STATIC_MODELS) + elif "omni" in name.lower(): + models.append(name) + cls._models_cache = models or (list(VEO_STATIC_MODELS) + list(OMNI_STATIC_MODELS)) except Exception: - cls._models_cache = list(STATIC_MODELS) + cls._models_cache = list(VEO_STATIC_MODELS) + list(OMNI_STATIC_MODELS) 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( @@ -61,25 +120,32 @@ class GeminiModel(Model): ) @classmethod - def video_settings_inputs(cls): - inputs = super().video_settings_inputs() - inputs.append( + def 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(), + ] + return [ + VEO_VIDEO_SETTINGS.size_input(), + VEO_VIDEO_SETTINGS.resolution_input(), + VEO_VIDEO_SETTINGS.seconds_input(), io.Combo.Input( "person_generation", options=["allow_all", "allow_adult"], default="allow_all", tooltip="Whether people may be generated in the video.", - ) - ) - inputs.append(cls._negative_prompt_input()) - return inputs + ), + cls._negative_prompt_input(), + ] @classmethod - def image_to_video_settings_inputs(cls): + def image_to_video_settings_inputs(cls, model_id: str | None = None): return [ - cls.VIDEO_SETTINGS.size_input(), - cls.VIDEO_SETTINGS.resolution_input(), - cls.VIDEO_SETTINGS.seconds_input(), + VEO_VIDEO_SETTINGS.size_input(), + VEO_VIDEO_SETTINGS.resolution_input(), + VEO_VIDEO_SETTINGS.seconds_input(), reference_image_input(), io.Combo.Input( "person_generation", @@ -92,9 +158,17 @@ class GeminiModel(Model): @classmethod def video_kwargs(cls, model_id: str, settings: dict) -> dict: + if _is_omni(model_id): + return { + "response_format": { + "type": "video", + "aspect_ratio": settings["size"], + "resolution": OMNI_RESOLUTION_LABELS[settings["resolution"]], + } + } kwargs = { "aspectRatio": settings["size"], - "resolution": RESOLUTION_MAP[settings["resolution"]], + "resolution": VEO_RESOLUTION_MAP[settings["resolution"]], "durationSeconds": settings["seconds"], "personGeneration": settings["person_generation"], } @@ -110,3 +184,20 @@ class GeminiModel(Model): kwargs = cls.video_kwargs(model_id, settings) kwargs["input_reference"] = cls.image_tensor_to_bytesio(reference_image) return kwargs + + @classmethod + def extract_video_bytes(cls, response) -> bytes: + for step in getattr(response, "steps", None) or []: + if step.get("type") != "model_output": + continue + for item in step.get("content") or []: + if item.get("type") != "video": + continue + data = item.get("data") + if data: + return base64.b64decode(data) + raise RuntimeError( + "Gemini Omni returned the video by reference (uri) instead of inline " + "base64 data; downloading a uri-delivered video isn't supported yet." + ) + raise RuntimeError("Gemini Omni response contained no video output.") diff --git a/src/models/openai.py b/src/models/openai.py index 7907772..f30b4b2 100644 --- a/src/models/openai.py +++ b/src/models/openai.py @@ -1,6 +1,6 @@ -from src.models.base import Model -from src.settings.base import Resolutions, Sizes -from src.settings.video import VideoSettings, dimensions_string +from .base import Model +from ..settings.base import Resolutions, Sizes +from ..settings.video import VideoSettings, dimensions_string class _SoraModel(Model): diff --git a/src/models/runwayml.py b/src/models/runwayml.py index 934a04d..2d0eae9 100644 --- a/src/models/runwayml.py +++ b/src/models/runwayml.py @@ -1,8 +1,8 @@ from comfy_api.latest import io -from src.models.base import Model -from src.settings.base import Resolutions, Sizes -from src.settings.video import VideoSettings, dimensions_string +from .base import Model +from ..settings.base import Resolutions, Sizes +from ..settings.video import VideoSettings, dimensions_string class RunwayMLModel(Model): @@ -20,7 +20,7 @@ class RunwayMLModel(Model): return ["gen3a_turbo", "gen4_turbo", "gen4_aleph", "gen4_image", "gen4_image_turbo"] @classmethod - def video_settings_inputs(cls): + def video_settings_inputs(cls, model_id: str | None = None): return [ cls.VIDEO_SETTINGS.size_input(), cls.VIDEO_SETTINGS.resolution_input(), diff --git a/src/nodes.py b/src/nodes.py index 68477bb..4836307 100644 --- a/src/nodes.py +++ b/src/nodes.py @@ -3,11 +3,12 @@ from io import BytesIO from comfy_api.latest import ComfyAPI, InputImpl, io from litellm import avideo_content, avideo_generation, avideo_status +from litellm.interactions import acreate as acreate_interaction -from src.models.azure import AzureModel -from src.models.gemini import GeminiModel -from src.models.openai import OpenAIModel -from src.models.runwayml import RunwayMLModel +from .models.azure import AzureModel +from .models.gemini import GeminiModel +from .models.openai import OpenAIModel +from .models.runwayml import RunwayMLModel api = ComfyAPI() @@ -126,6 +127,26 @@ class APIVideoGenerate(io.ComfyNode): extra_kwargs = provider_cls.image_to_video_kwargs(model_id, settings) model = f"{provider_name}/{model_id}" + if provider_cls.uses_interactions_api(model_id): + # The Interactions API is a single request/response call with no + # status-polling endpoint, so there's no real progress fraction to + # 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. + request = asyncio.ensure_future( + acreate_interaction(model=model, input=prompt, timeout=timeout, **auth_kwargs, **extra_kwargs) + ) + elapsed = 0.0 + await api.execution.set_progress(0.0, timeout, node_id=node_id) + while not request.done(): + await asyncio.wait([request], timeout=poll_interval) + elapsed = min(elapsed + poll_interval, timeout) + await api.execution.set_progress(elapsed, timeout, node_id=node_id) + response = await request + await api.execution.set_progress(timeout, timeout, node_id=node_id) + video_bytes = provider_cls.extract_video_bytes(response) + return io.NodeOutput(InputImpl.VideoFromFile(BytesIO(video_bytes))) + op = await avideo_generation( model=model, prompt=prompt, timeout=timeout, **auth_kwargs, **extra_kwargs ) diff --git a/src/settings/base.py b/src/settings/base.py index d1ac3c7..dfe5d9d 100644 --- a/src/settings/base.py +++ b/src/settings/base.py @@ -25,6 +25,7 @@ class Resolutions(str, Enum): once; every model's RESOLUTIONS list references these members instead of a raw string.""" + _360 = "360" _480 = "480" _720 = "720" _1080 = "1080" diff --git a/src/settings/image.py b/src/settings/image.py index d368c68..9812a7e 100644 --- a/src/settings/image.py +++ b/src/settings/image.py @@ -1,6 +1,6 @@ from comfy_api.latest import io -from src.settings.base import Resolutions, Settings, Sizes +from .base import Resolutions, Settings, Sizes IMAGE_SIZES = [ Sizes._21_9, diff --git a/src/settings/video.py b/src/settings/video.py index f3e09e0..eb1997a 100644 --- a/src/settings/video.py +++ b/src/settings/video.py @@ -1,6 +1,6 @@ from comfy_api.latest import io -from src.settings.base import Resolutions, Settings, Sizes +from .base import Resolutions, Settings, Sizes VIDEO_SIZES = [Sizes._21_9, Sizes._16_9, Sizes._4_3, Sizes._1_1, Sizes._3_4, Sizes._9_16] VIDEO_RESOLUTIONS = [Resolutions._1K, Resolutions._2K, Resolutions._4K] diff --git a/tests/test_gemini_list_models.py b/tests/test_gemini_list_models.py new file mode 100644 index 0000000..3273581 --- /dev/null +++ b/tests/test_gemini_list_models.py @@ -0,0 +1,88 @@ +import sys +import types + +import pytest + +from src.models.gemini import GeminiModel + + +class FakeGenaiModel: + def __init__(self, name, supported_actions=None): + self.name = name + self.supported_actions = supported_actions + + +def _install_fake_genai(models): + fake_models = types.SimpleNamespace(list=lambda: models) + fake_client_instance = types.SimpleNamespace(models=fake_models) + fake_genai_module = types.SimpleNamespace(Client=lambda: fake_client_instance) + google_module = types.ModuleType("google") + google_module.genai = fake_genai_module + sys.modules["google"] = google_module + sys.modules["google.genai"] = fake_genai_module + + +@pytest.fixture(autouse=True) +def _reset_cache(): + GeminiModel._models_cache = None + yield + GeminiModel._models_cache = None + sys.modules.pop("google", None) + sys.modules.pop("google.genai", None) + + +def test_excludes_veo_model_missing_predict_long_running(): + _install_fake_genai( + [ + FakeGenaiModel("models/veo-3.1-fast-generate-001", supported_actions=["generateContent"]), + FakeGenaiModel("models/veo-3.1-fast-generate-preview", supported_actions=["predictLongRunning"]), + ] + ) + + models = GeminiModel.list_models() + + assert "veo-3.1-fast-generate-001" not in models + assert "veo-3.1-fast-generate-preview" in models + + +def test_keeps_veo_model_with_unknown_supported_actions(): + _install_fake_genai( + [ + FakeGenaiModel("models/veo-3.1-generate-preview", supported_actions=None), + ] + ) + + models = GeminiModel.list_models() + + assert "veo-3.1-generate-preview" in models + + +def test_omni_models_not_filtered_by_supported_actions(): + _install_fake_genai( + [ + FakeGenaiModel("models/gemini-omni-flash-preview", supported_actions=["generateContent"]), + ] + ) + + models = GeminiModel.list_models() + + assert "gemini-omni-flash-preview" in models + + +def test_falls_back_to_static_models_when_genai_list_fails(): + def _raise(): + raise RuntimeError("no api key") + + fake_models = types.SimpleNamespace(list=_raise) + fake_client_instance = types.SimpleNamespace(models=fake_models) + fake_genai_module = types.SimpleNamespace(Client=lambda: fake_client_instance) + google_module = types.ModuleType("google") + google_module.genai = fake_genai_module + sys.modules["google"] = google_module + sys.modules["google.genai"] = fake_genai_module + + models = GeminiModel.list_models() + + assert "veo-3.1-fast-generate-001" not in models + assert "veo-3.1-fast-generate-preview" in models + assert "gemini-omni-1.1-flash" in models diff --git a/tests/test_nodes.py b/tests/test_nodes.py index 2472c98..2847404 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -1,4 +1,5 @@ import asyncio +import base64 import types from unittest.mock import AsyncMock, patch @@ -197,12 +198,28 @@ def test_schema_builds_nested_source_provider_model_combo(): field_names = {i.name for i in model_option.inputs} assert {"size", "resolution", "seconds"} <= field_names + gemini_text_option = next(o for o in provider_input.options if o.key == "gemini") + gemini_text_model_input = gemini_text_option.inputs[0] + gemini_text_model_ids = {o.key for o in gemini_text_model_input.options} + assert "veo-3.1-fast-generate-preview" in gemini_text_model_ids + assert "gemini-omni-flash-preview" in gemini_text_model_ids + + veo_option = next(o for o in gemini_text_model_input.options if o.key == "veo-3.1-fast-generate-preview") + veo_field_names = {i.name for i in veo_option.inputs} + assert "person_generation" in veo_field_names + + omni_option = next(o for o in gemini_text_model_input.options if o.key == "gemini-omni-flash-preview") + omni_field_names = {i.name for i in omni_option.inputs} + assert "person_generation" not in omni_field_names + image_provider_input = image_option.inputs[0] assert image_provider_input.name == "provider" assert [o.key for o in image_provider_input.options] == ["gemini"] 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 @@ -282,6 +299,88 @@ def test_gemini_image_to_video_requires_reference_image(): asyncio.run(run()) +class FakeInteractionResponse: + def __init__(self, steps): + self.steps = steps + + +def _gemini_omni_source(model="gemini-omni-flash-preview", **overrides): + settings = {"model": model, "size": "16:9", "resolution": "720", "seconds": 5} + settings.update(overrides) + return {"source": "text", "provider": {"provider": "gemini", "model": settings}} + + +def test_gemini_omni_uses_interactions_api_not_polling(): + async def run(): + interaction_mock = AsyncMock( + return_value=FakeInteractionResponse( + steps=[ + {"type": "user_input", "content": [{"type": "text", "text": "hi"}]}, + { + "type": "model_output", + "content": [ + { + "type": "video", + "data": base64.b64encode(b"fake-mp4-bytes").decode(), + "mime_type": "video/mp4", + } + ], + }, + ] + ) + ) + gen_mock = AsyncMock() + + with patch.object(nodes, "acreate_interaction", interaction_mock), patch.object( + nodes, "avideo_generation", gen_mock + ): + result = await nodes.APIVideoGenerate.execute( + prompt="a marble rolling down a track", + source=_gemini_omni_source(), + ) + + gen_mock.assert_not_awaited() + interaction_mock.assert_awaited_once() + kwargs = interaction_mock.call_args.kwargs + assert kwargs["model"] == "gemini/gemini-omni-flash-preview" + assert kwargs["input"] == "a marble rolling down a track" + assert kwargs["response_format"] == { + "type": "video", + "aspect_ratio": "16:9", + "resolution": "720p", + } + assert "api_key" not in kwargs + + video = result.args[0] + assert isinstance(video, nodes.InputImpl.VideoFromFile) + assert video.file_obj.getvalue() == b"fake-mp4-bytes" + + asyncio.run(run()) + + +def test_gemini_omni_uri_delivery_raises_clear_error(): + async def run(): + interaction_mock = AsyncMock( + return_value=FakeInteractionResponse( + steps=[ + { + "type": "model_output", + "content": [{"type": "video", "uri": "https://example.com/video.mp4"}], + } + ] + ) + ) + + with patch.object(nodes, "acreate_interaction", interaction_mock): + with pytest.raises(RuntimeError, match="uri"): + await nodes.APIVideoGenerate.execute( + prompt="anything", + source=_gemini_omni_source(), + ) + + asyncio.run(run()) + + def test_azure_provider_string_uses_azure_prefix(): async def run(): gen_mock = AsyncMock(return_value=FakeVideoObject("completed"))