From d2f1f4950fc99d8d90e86cdce42677c33080ab6e Mon Sep 17 00:00:00 2001 From: silveroxides Date: Sun, 30 Aug 2026 03:11:18 +0200 Subject: [PATCH] feat(video): add configurable Gemini input --- README.md | 8 +- __init__.py | 7 +- gemini_nodes.py | 111 ++++++++++++++++++++++++--- tests/test_gemini_nodes.py | 151 ++++++++++++++++++++++++++++++++++--- tests/test_groups.toml | 1 + 5 files changed, 253 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index 4e0f5bb..f5de778 100644 --- a/README.md +++ b/README.md @@ -56,8 +56,9 @@ Use the `Expanded Gemini Text/Image` node for text generation: - Connect the API key configuration node to the `config` input. - Enter your prompt text in `prompt`. - Adjust generation parameters (temperature, top_p, top_k, etc.). -- If needed, connect an image to the `input_image` input for image understanding. -- Enable `use_cache` on the configuration node to reuse unchanged prompt, system instruction, and image context through Gemini after the local result cache misses. +- If needed, connect images through `image_inputs`. For video, pass a native ComfyUI video through `Configure Gemini Video Input`, then connect its output to `video`. +- Configured video is normalized to MP4/H.264 and sent with embedded audio at the selected sampling FPS. The normalized video must be smaller than 100 MB. +- Enable `use_cache` on the configuration node to reuse unchanged prompt, system instruction, image, and video context through Gemini after the local result cache misses. ### 3. Proxy Settings @@ -91,7 +92,8 @@ If you are in China or other regions requiring a proxy: #### Optional Parameters -- `input_image`: Input image (for image understanding) +- `video`: Optional `GEMINI_VIDEO_CONFIG` input. Its video is sent before images and prompt with embedded audio. +- `image_inputs`: Optional ordered image autogrow inputs for image understanding. - `use_proxy`: Whether to use a proxy (True/False) - `proxy_host`: Proxy host address - `proxy_port`: Proxy port diff --git a/__init__.py b/__init__.py index 302d91b..c001369 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,5 @@ from comfy_api.latest import ComfyExtension, io -from .gemini_nodes import GetKeyAPI, SSL_GeminiAPIKeyConfig, SSL_GeminiTextPrompt +from .gemini_nodes import GetKeyAPI, SSL_GeminiAPIKeyConfig, SSL_GeminiVideoConfig, SSL_GeminiTextPrompt class GeminiExtension(ComfyExtension): @@ -9,8 +9,9 @@ class GeminiExtension(ComfyExtension): async def get_node_list(self) -> list[type[io.ComfyNode]]: return [ GetKeyAPI, - SSL_GeminiAPIKeyConfig, - SSL_GeminiTextPrompt, + SSL_GeminiAPIKeyConfig, + SSL_GeminiVideoConfig, + SSL_GeminiTextPrompt, ] diff --git a/gemini_nodes.py b/gemini_nodes.py index a1e8f88..043cbaf 100644 --- a/gemini_nodes.py +++ b/gemini_nodes.py @@ -8,7 +8,7 @@ import cv2 from PIL import Image from io import BytesIO import folder_paths # type: ignore[reportMissingImports] -from comfy_api.latest import ComfyExtension, UI, IO # type: ignore[reportMissingImports] +from comfy_api.latest import ComfyExtension, UI, IO, Types # type: ignore[reportMissingImports] from google import genai from google.genai import errors, types import time @@ -145,14 +145,39 @@ class SSL_GeminiAPIKeyConfig(IO.ComfyNode): +class SSL_GeminiVideoConfig(IO.ComfyNode): + GeminiVideoConfig = IO.Custom("GEMINI_VIDEO_CONFIG") + + @classmethod + def define_schema(cls) -> IO.Schema: + return IO.Schema( + node_id="SSL_GeminiVideoConfig", + display_name="Configure Gemini Video Input", + category="API/Gemini", + inputs=[ + IO.Video.Input("video"), + IO.Int.Input("fps", default=1, min=1, max=24, step=1, tooltip="Frames per second sampled by Gemini for video understanding."), + ], + outputs=[cls.GeminiVideoConfig.Output("video_config")], + ) + + @classmethod + def execute(cls, video: IO.Video.Type, fps: int) -> IO.NodeOutput: + return IO.NodeOutput({"video": video, "fps": int(fps)}) + + class SSL_GeminiTextPrompt(IO.ComfyNode): GemConfig = IO.Custom("GEMINI_CONFIG") + GeminiVideoConfig = IO.Custom("GEMINI_VIDEO_CONFIG") _cache: dict = {} _seed_map_cache: dict = {} # Maps (input_seed, fingerprint_without_seed) -> successful_gemini_seed _client_cache: dict = {} # Maps client_key tuple -> genai.Client instance _context_cache: dict = {} _context_cache_lock = threading.Lock() GEMINI_3_7_FLASH = "gemini-3.7-flash" + VIDEO_MIME_TYPE = "video/mp4" + VIDEO_INLINE_LIMIT_BYTES = 100 * 1024 * 1024 + VERTEX_CACHE_INLINE_LIMIT_BYTES = 10 * 1024 * 1024 # Placeholder only. Keep this out of the model combo until Google's hosted-model # documentation confirms the final ID and gemini-3.6-flash-equivalent capabilities @@ -216,6 +241,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): IO.String.Input("retry_pattern", default="", optional=True, multiline=False, tooltip="Regex pattern to match in response text. If matched, retry with new seed. Leave empty to disable."), IO.Int.Input("max_retries", default=3, min=0, max=10, step=1, tooltip="Maximum number of retry attempts when pattern matches. 0 disables retry."), IO.String.Input("timeout_fallback_text", default="", optional=True, multiline=True, tooltip="Text returned when the Gemini request times out. Leave empty to return the standard timeout message."), + cls.GeminiVideoConfig.Input("video", optional=True, tooltip="Optional configured Gemini video input with embedded audio and sampling FPS."), IO.Autogrow.Input( "image_inputs", template=image_template, @@ -320,6 +346,22 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): ) return frames, batch_counts + @classmethod + def _serialize_video(cls, video: IO.Video.Type | None): + if video is None: + return None, None + + buffer = BytesIO() + video.save_to( + buffer, + format=Types.VideoContainer.MP4, + codec=Types.VideoCodec.H264, + ) + video_bytes = buffer.getvalue() + if len(video_bytes) >= cls.VIDEO_INLINE_LIMIT_BYTES: + raise ValueError("Gemini video input must be smaller than 100 MB after MP4/H.264 normalization.") + return video_bytes, cls.VIDEO_MIME_TYPE + @classmethod def _resolve_gemini_3_7_thinking_level(cls, thinking_level): if thinking_level is None or thinking_level == "None": @@ -339,6 +381,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): @classmethod def _compute_fingerprint_and_check_cache(cls, config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens, include_images, aspect_ratio, bypass_mode, thinking_budget, use_seed, seed, + video_hash=None, video_mime_type=None, video_fps=None, image_inputs: IO.Autogrow.Type | None = None, use_proxy=False, proxy_host="127.0.0.1", proxy_port=7890, timeout=30, include_thoughts=False, thinking_level=None, media_resolution=None, @@ -418,6 +461,9 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): bool(use_cache), int(cache_ttl_minutes) if use_cache else 0, int(cache_seed) if use_cache and use_seed else 0, + video_hash, + video_mime_type, + video_fps, image_hashes, bool(use_proxy), str(proxy_host), @@ -505,12 +551,15 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): @staticmethod def _build_context_cache_key(client_key, model, system_instruction, prompt, - image_hashes, media_resolution): + video_hash, video_mime_type, video_fps, image_hashes, media_resolution): return ( client_key, model, system_instruction, prompt, + video_hash, + video_mime_type, + video_fps, tuple(image_hashes), str(media_resolution), ) @@ -665,15 +714,26 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): use_proxy=False, proxy_host="127.0.0.1", proxy_port=7890, use_seed=False, seed=0, timeout=30, include_thoughts=False, thinking_level=None, media_resolution=None, retry_pattern="", max_retries=3, timeout_fallback_text="", + video: GeminiVideoConfig.Type | None = None, image_inputs: IO.Autogrow.Type | None = None) -> IO.NodeOutput: print(f"[INFO] SSL_GeminiTextPrompt execute called, model: {model}") use_cache = bool(config.get("use_cache", False)) cache_ttl_minutes = int(config.get("cache_ttl_minutes", 60)) cache_seed = int(config.get("cache_seed", 0)) + video_input = video.get("video") if video is not None else None + video_fps = int(video.get("fps", 1)) if video is not None else None + try: + video_bytes, video_mime_type = cls._serialize_video(video_input) + except Exception as e: + print(f"[ERROR] Error processing input video: {e}") + output_seed = cache_seed if use_cache and use_seed else seed if use_seed else 0 + return IO.NodeOutput(f"Error processing input video: {e}", cls.generate_empty_image(), output_seed) + video_hash = hashlib.sha256(video_bytes).hexdigest() if video_bytes is not None else None fingerprint, cached = cls._compute_fingerprint_and_check_cache( config, prompt, system_instruction, model, temperature, top_p, top_k, max_output_tokens, include_images, aspect_ratio, bypass_mode, thinking_budget, use_seed, seed, + video_hash, video_mime_type, video_fps, image_inputs, use_proxy, proxy_host, proxy_port, timeout, include_thoughts, thinking_level, media_resolution, @@ -837,9 +897,27 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): print(f"[ERROR] Gemini client initialization failed: {str(e)}") return IO.NodeOutput(f"Gemini client initialization failed: {str(e)}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0) - # Prepare contents (images + prompt) + # Prepare contents (video + images + prompt) image_frames, image_batch_counts = cls._ordered_image_frames(image_inputs) context_image_hashes = [] + media_parts = [] + context_inline_bytes = ( + (len(video_bytes) if video_bytes is not None else 0) + + len(padded_prompt.encode("utf-8")) + + len(padded_system_instruction.encode("utf-8")) + ) + + if video_bytes is not None: + if model in cls.MEDIA_RES_MODELS and media_resolution is not None and media_resolution != "unspecified": + video_part = types.Part.from_bytes( + data=video_bytes, + mime_type=video_mime_type, + media_resolution=f"MEDIA_RESOLUTION_{media_resolution.upper()}", + ) + else: + video_part = types.Part.from_bytes(data=video_bytes, mime_type=video_mime_type) + video_part.video_metadata = types.VideoMetadata(fps=video_fps) + media_parts.append(video_part) if image_frames: try: @@ -851,6 +929,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): img_byte_arr = BytesIO() pil_img.save(img_byte_arr, format='PNG') img_bytes = img_byte_arr.getvalue() + context_inline_bytes += len(img_bytes) context_image_hashes.append(hashlib.sha256(img_bytes).hexdigest()) if model in cls.MEDIA_RES_MODELS and media_resolution is not None and media_resolution != "unspecified": img_part = types.Part.from_bytes( @@ -872,14 +951,16 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): f"[INFO] Prepared Gemini image inputs: {batch_summary}, " f"total={len(img_parts)}" ) - contents = [ - types.UserContent( - parts=img_parts + [types.Part.from_text(text=padded_prompt)] - ) - ] + media_parts.extend(img_parts) except Exception as e: print(f"[ERROR] Error processing input image: {str(e)}") return IO.NodeOutput(f"Error processing input image: {str(e)}", cls.generate_empty_image(), actual_seed if actual_seed is not None else 0) + if media_parts: + contents = [ + types.UserContent( + parts=media_parts + [types.Part.from_text(text=padded_prompt)] + ) + ] else: contents = padded_prompt @@ -891,10 +972,22 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): model, padded_system_instruction, padded_prompt, + video_hash, + video_mime_type, + video_fps, context_image_hashes, media_resolution, ) + skip_context_cache = ( + use_cache + and video_bytes is not None + and (config.get("use_vertexai_env", False) or config.get("vertexai_express", False)) + and context_inline_bytes > cls.VERTEX_CACHE_INLINE_LIMIT_BYTES + ) + if skip_context_cache: + print("[WARNING] Vertex context cache skipped because inline cached media exceeds 10 MB; using full request") + generate_content_config = cls._build_generate_content_config( model=model, temperature=temperature, @@ -928,7 +1021,7 @@ class SSL_GeminiTextPrompt(IO.ComfyNode): result_queue: "queue.Queue[Tuple[str, Any]]" = queue.Queue() def generate_once(): - if not use_cache: + if not use_cache or skip_context_cache: return client.models.generate_content( model=model, contents=contents, diff --git a/tests/test_gemini_nodes.py b/tests/test_gemini_nodes.py index 3e1160a..851e768 100644 --- a/tests/test_gemini_nodes.py +++ b/tests/test_gemini_nodes.py @@ -59,13 +59,41 @@ def _install_success_client(monkeypatch, captured): gemini_nodes.SSL_GeminiTextPrompt._client_cache.clear() +class FakeVideo: + def __init__(self, data): + self.data = data + self.save_calls = [] + + def save_to(self, buffer, format, codec): + self.save_calls.append((format, codec)) + buffer.write(self.data) + + +def _video_config(data, fps=1): + return {"video": FakeVideo(data), "fps": fps} + + def test_image_autogrow_is_last_input(): schema = gemini_nodes.SSL_GeminiTextPrompt.define_schema() - assert schema.inputs[-2].id == "timeout_fallback_text" + assert schema.inputs[-2].id == "video" + assert schema.inputs[-2].optional is True assert schema.inputs[-1].id == "image_inputs" +def test_video_config_exposes_native_video_and_integer_fps(): + schema = gemini_nodes.SSL_GeminiVideoConfig.define_schema() + fps = next(input_ for input_ in schema.inputs if input_.id == "fps") + + assert [input_.id for input_ in schema.inputs] == ["video", "fps"] + assert fps.default == 1 + assert fps.min == 1 + assert fps.max == 24 + native_video = object() + output = gemini_nodes.SSL_GeminiVideoConfig.execute(native_video, 2) + assert output[0] == {"video": native_video, "fps": 2} + + def test_context_cache_reuses_full_input_after_local_fingerprint_miss(monkeypatch): captured = {"creates": [], "generations": []} @@ -112,6 +140,7 @@ def test_context_cache_reuses_full_input_after_local_fingerprint_miss(monkeypatc kwargs.update( use_seed=True, seed=1, + video=_video_config(b"normalized-video", fps=2), image_inputs={"image_1": torch.full((1, 2, 2, 3), 0.25)}, ) @@ -125,8 +154,11 @@ def test_context_cache_reuses_full_input_after_local_fingerprint_miss(monkeypatc create_config = captured["creates"][0]["config"] assert create_config.system_instruction == "test system instruction" assert create_config.ttl == "900s" + assert create_config.contents[0].parts[0].inline_data.data == b"normalized-video" + assert create_config.contents[0].parts[0].inline_data.mime_type == "video/mp4" + assert create_config.contents[0].parts[0].video_metadata.fps == 2 assert create_config.contents[0].parts[-1].text == "test prompt" - assert len(create_config.contents[0].parts) == 2 + assert len(create_config.contents[0].parts) == 3 assert len(captured["generations"]) == 2 assert all(call["contents"] == " " for call in captured["generations"]) assert all(call["config"].cached_content == "cachedContents/test" for call in captured["generations"]) @@ -284,18 +316,117 @@ def test_context_cache_ttl_change_updates_existing_resource(): def test_context_cache_identity_covers_every_cached_input(): build_key = gemini_nodes.SSL_GeminiTextPrompt._build_context_cache_key - base = build_key(("standard", "key-a"), "model-a", "system-a", "prompt-a", ["image-a"], "low") + base = build_key(("standard", "key-a"), "model-a", "system-a", "prompt-a", "video-a", "video/mp4", 1, ["image-a"], "low") variants = { - build_key(("standard", "key-b"), "model-a", "system-a", "prompt-a", ["image-a"], "low"), - build_key(("standard", "key-a"), "model-b", "system-a", "prompt-a", ["image-a"], "low"), - build_key(("standard", "key-a"), "model-a", "system-b", "prompt-a", ["image-a"], "low"), - build_key(("standard", "key-a"), "model-a", "system-a", "prompt-b", ["image-a"], "low"), - build_key(("standard", "key-a"), "model-a", "system-a", "prompt-a", ["image-b"], "low"), - build_key(("standard", "key-a"), "model-a", "system-a", "prompt-a", ["image-a"], "high"), + build_key(("standard", "key-b"), "model-a", "system-a", "prompt-a", "video-a", "video/mp4", 1, ["image-a"], "low"), + build_key(("standard", "key-a"), "model-b", "system-a", "prompt-a", "video-a", "video/mp4", 1, ["image-a"], "low"), + build_key(("standard", "key-a"), "model-a", "system-b", "prompt-a", "video-a", "video/mp4", 1, ["image-a"], "low"), + build_key(("standard", "key-a"), "model-a", "system-a", "prompt-b", "video-a", "video/mp4", 1, ["image-a"], "low"), + build_key(("standard", "key-a"), "model-a", "system-a", "prompt-a", "video-b", "video/mp4", 1, ["image-a"], "low"), + build_key(("standard", "key-a"), "model-a", "system-a", "prompt-a", "video-a", "video/webm", 1, ["image-a"], "low"), + build_key(("standard", "key-a"), "model-a", "system-a", "prompt-a", "video-a", "video/mp4", 2, ["image-a"], "low"), + build_key(("standard", "key-a"), "model-a", "system-a", "prompt-a", "video-a", "video/mp4", 1, ["image-b"], "low"), + build_key(("standard", "key-a"), "model-a", "system-a", "prompt-a", "video-a", "video/mp4", 1, ["image-a"], "high"), } assert base not in variants - assert len(variants) == 6 + assert len(variants) == 9 + + +def test_video_is_normalized_sent_first_and_locally_fingerprinted(monkeypatch): + captured = {"generations": []} + + class FakeModels: + @staticmethod + def generate_content(**kwargs): + captured["generations"].append(kwargs) + part = SimpleNamespace(text="ok", inline_data=None) + return SimpleNamespace(candidates=[SimpleNamespace(content=SimpleNamespace(parts=[part]))]) + + class FakeClient: + def __init__(self, **kwargs): + self.models = FakeModels() + + monkeypatch.setattr(gemini_nodes.genai, "Client", FakeClient) + gemini_nodes.SSL_GeminiTextPrompt._cache.clear() + gemini_nodes.SSL_GeminiTextPrompt._client_cache.clear() + kwargs = _execute_kwargs(_config(api_key="test-key")) + kwargs.update( + model="gemini-3-flash-preview", + media_resolution="low", + use_seed=True, + seed=7, + video=_video_config(b"video-a", fps=2), + ) + + first = gemini_nodes.SSL_GeminiTextPrompt.execute(**kwargs) + kwargs["video"] = _video_config(b"video-a", fps=2) + second = gemini_nodes.SSL_GeminiTextPrompt.execute(**kwargs) + kwargs["video"] = _video_config(b"video-b", fps=2) + third = gemini_nodes.SSL_GeminiTextPrompt.execute(**kwargs) + + assert first[0] == second[0] == third[0] == "ok" + assert len(captured["generations"]) == 2 + parts = captured["generations"][0]["contents"][0].parts + assert parts[0].inline_data.data == b"video-a" + assert parts[0].inline_data.mime_type == "video/mp4" + assert parts[0].media_resolution.level == "MEDIA_RESOLUTION_LOW" + assert parts[-1].text == "test prompt" + assert kwargs["video"]["video"].save_calls == [ + (gemini_nodes.Types.VideoContainer.MP4, gemini_nodes.Types.VideoCodec.H264) + ] + + +def test_video_at_inline_limit_returns_clear_error(monkeypatch): + monkeypatch.setattr(gemini_nodes.SSL_GeminiTextPrompt, "VIDEO_INLINE_LIMIT_BYTES", 4) + output = gemini_nodes.SSL_GeminiTextPrompt.execute( + **_execute_kwargs(_config()), + video=_video_config(b"1234"), + ) + + assert output[0] == "Error processing input video: Gemini video input must be smaller than 100 MB after MP4/H.264 normalization." + + +def test_vertex_video_over_cache_blob_limit_uses_full_request(monkeypatch): + captured = {"cache_creates": 0, "generations": []} + + class FakeCaches: + @staticmethod + def create(**kwargs): + captured["cache_creates"] += 1 + raise AssertionError("Vertex cache creation must be skipped") + + class FakeModels: + @staticmethod + def generate_content(**kwargs): + captured["generations"].append(kwargs) + part = SimpleNamespace(text="ok", inline_data=None) + return SimpleNamespace(candidates=[SimpleNamespace(content=SimpleNamespace(parts=[part]))]) + + class FakeClient: + def __init__(self, **kwargs): + self.caches = FakeCaches() + self.models = FakeModels() + + monkeypatch.setattr(gemini_nodes.genai, "Client", FakeClient) + monkeypatch.setattr(gemini_nodes.SSL_GeminiTextPrompt, "VERTEX_CACHE_INLINE_LIMIT_BYTES", 1) + gemini_nodes.SSL_GeminiTextPrompt._cache.clear() + gemini_nodes.SSL_GeminiTextPrompt._client_cache.clear() + gemini_nodes.SSL_GeminiTextPrompt._context_cache.clear() + config = _config(api_key="test-key", vertexai_express=True, use_cache=True, cache_seed=11) + + output = gemini_nodes.SSL_GeminiTextPrompt.execute( + **_execute_kwargs(config), + use_seed=True, + seed=2, + video=_video_config(b"12"), + ) + + assert output[0] == "ok" + assert captured["cache_creates"] == 0 + assert len(captured["generations"]) == 1 + assert captured["generations"][0]["contents"] != " " + assert captured["generations"][0]["config"].seed == 11 def test_vertex_config_socket_ids_remain_stable(): diff --git a/tests/test_groups.toml b/tests/test_groups.toml index f71a03b..93867e1 100644 --- a/tests/test_groups.toml +++ b/tests/test_groups.toml @@ -1,6 +1,7 @@ # Gemini node schema, request construction, caching, and authentication behavior. [groups.gemini_node] paths = [ + "__init__.py", "gemini_nodes.py", ] python_tests = [