From 6691b570a72374aec5b914f4c40d20bef26e661c Mon Sep 17 00:00:00 2001 From: Andy Ratsirarson Date: Wed, 15 Jul 2026 22:42:27 -0700 Subject: [PATCH] Allow input image socket node --- README.md | 2 +- src/models/base.py | 25 +++++++++++++++++++--- src/models/gemini.py | 8 +++---- src/settings/video.py | 20 ++++++------------ tests/conftest.py | 49 +++++++------------------------------------ tests/test_nodes.py | 39 +++++++++++++++++++++++++++++++--- 6 files changed, 77 insertions(+), 66 deletions(-) diff --git a/README.md b/README.md index ec2a1bf..e379291 100644 --- a/README.md +++ b/README.md @@ -81,7 +81,7 @@ 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`. -4. Set `reference_image`. This is a file picker + upload widget that works exactly like ComfyUI's built-in **Load Image** node: choose an image already in ComfyUI's input folder from the dropdown, or click upload to add a new one. You can also drag an image straight from the results gallery onto the widget — ComfyUI copies it into the input folder and selects it. (It's a widget on the node, not a socket you wire from another node's output.) +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`. diff --git a/src/models/base.py b/src/models/base.py index adf8746..31638a7 100644 --- a/src/models/base.py +++ b/src/models/base.py @@ -45,11 +45,30 @@ class Model: @classmethod def image_to_video_kwargs(cls, model_id: str, settings: dict) -> dict: - """Translate the generic settings dict (including the resolved reference - image path) into the litellm avideo_generation kwargs this provider's - API expects for image-to-video.""" + """Translate the generic settings dict (including the reference image + tensor) into the litellm avideo_generation kwargs this provider's API + expects for image-to-video.""" raise NotImplementedError(f"image_to_video_kwargs 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 + a PNG-encoded BytesIO, for providers whose API needs image bytes/a + file-like object (e.g. litellm's avideo_generation input_reference). + 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 + 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 + @classmethod def model_combo_option(cls, mode="text"): """Build this provider's DynamicCombo.Option: a nested 'model' combo where diff --git a/src/models/gemini.py b/src/models/gemini.py index e88199a..3c75c49 100644 --- a/src/models/gemini.py +++ b/src/models/gemini.py @@ -1,4 +1,3 @@ -import folder_paths from comfy_api.latest import io from src.models.base import Model @@ -105,8 +104,9 @@ class GeminiModel(Model): @classmethod def image_to_video_kwargs(cls, model_id: str, settings: dict) -> dict: + reference_image = settings.get("reference_image") + 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"] = folder_paths.get_annotated_filepath( - settings["reference_image"] - ) + kwargs["input_reference"] = cls.image_tensor_to_bytesio(reference_image) return kwargs diff --git a/src/settings/video.py b/src/settings/video.py index 8ae218c..f3e09e0 100644 --- a/src/settings/video.py +++ b/src/settings/video.py @@ -1,6 +1,3 @@ -import os - -import folder_paths from comfy_api.latest import io from src.settings.base import Resolutions, Settings, Sizes @@ -29,19 +26,14 @@ def dimensions_string(size: str, resolution: str, separator: str) -> str: def reference_image_input(id="reference_image", optional=True): - """A file-picker + upload widget for image-to-video, matching ComfyUI's - built-in LoadImage/LoadVideo convention: lists existing files already in - the input directory and lets the user upload a new one. Yields a filename; - resolve it to a real path with folder_paths.get_annotated_filepath().""" - input_dir = folder_paths.get_input_directory() - files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] - files = folder_paths.filter_files_content_types(files, ["image"]) - return io.Combo.Input( + """An IMAGE socket for image-to-video, matching how every other ComfyUI + BYOK node (Recraft, OpenAI, BFL, Kling, ...) takes a reference image: + connect a Load Image node, or any other node's IMAGE output — including + an image generated earlier in the same workflow, no save-to-disk needed.""" + return io.Image.Input( id, - options=sorted(files), - upload=io.UploadType.image, optional=optional, - tooltip="Reference image to seed the video. Pick an existing upload or click to upload a new one.", + tooltip="Reference image to seed the video. Connect a Load Image node or any other node's IMAGE output.", ) diff --git a/tests/conftest.py b/tests/conftest.py index 8a446c0..7fb5f40 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,14 +1,11 @@ -"""Stubs the parts of `comfy_api.latest` (and `folder_paths`) our nodes touch. +"""Stubs the parts of `comfy_api.latest` our nodes touch. -Both only exist inside a running ComfyUI install, so unit tests fake just -enough of their surface (io.Schema/Input/Output/NodeOutput, ComfyAPI, -InputImpl.VideoFromFile, folder_paths' input-directory helpers) to import and -exercise src/nodes.py in isolation. +`comfy_api` only exists inside a running ComfyUI install, so unit tests fake +just enough of its surface (io.Schema/Input/Output/NodeOutput, ComfyAPI, +InputImpl.VideoFromFile) to import and exercise src/nodes.py in isolation. """ -import os import sys -import tempfile import types @@ -48,6 +45,9 @@ def _install_fake_comfy_api(): class Combo(_InputFactory): pass + class Image(_InputFactory): + pass + class _DynamicComboOption: def __init__(self, key, inputs): self.key = key @@ -68,12 +68,6 @@ def _install_fake_comfy_api(): class Video(_OutputFactory): pass - class UploadType: - image = "image_upload" - audio = "audio_upload" - video = "video_upload" - model = "file_upload" - class Hidden: unique_id = "unique_id" @@ -93,9 +87,9 @@ def _install_fake_comfy_api(): Float=Float, Int=Int, Combo=Combo, + Image=Image, DynamicCombo=DynamicCombo, Video=Video, - UploadType=UploadType, Hidden=Hidden, Schema=Schema, NodeOutput=NodeOutput, @@ -134,31 +128,4 @@ def _install_fake_comfy_api(): sys.modules["comfy_api.latest"] = latest -def _install_fake_folder_paths(): - if "folder_paths" in sys.modules: - return - - fake_input_dir = tempfile.mkdtemp(prefix="fake_comfy_input_") - - def get_input_directory(): - return fake_input_dir - - def filter_files_content_types(files, content_types): - return files - - def get_annotated_filepath(name, default_dir=None): - return os.path.join(default_dir or fake_input_dir, name) - - def exists_annotated_filepath(name): - return os.path.isfile(os.path.join(fake_input_dir, name)) - - fake_module = types.ModuleType("folder_paths") - fake_module.get_input_directory = get_input_directory - fake_module.filter_files_content_types = filter_files_content_types - fake_module.get_annotated_filepath = get_annotated_filepath - fake_module.exists_annotated_filepath = exists_annotated_filepath - sys.modules["folder_paths"] = fake_module - - _install_fake_comfy_api() -_install_fake_folder_paths() diff --git a/tests/test_nodes.py b/tests/test_nodes.py index 975d015..2472c98 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -218,6 +218,9 @@ def test_gemini_image_to_video_kwargs_mapping(): gen_mock = AsyncMock(return_value=FakeVideoObject("completed")) content_mock = AsyncMock(return_value=b"bytes") + fake_tensor = object() + fake_bytesio = object() + source = { "source": "image", "provider": { @@ -227,7 +230,7 @@ def test_gemini_image_to_video_kwargs_mapping(): "size": "9:16", "resolution": "2K", "seconds": 8, - "reference_image": "my_photo.png", + "reference_image": fake_tensor, "person_generation": "allow_adult", }, }, @@ -235,16 +238,46 @@ def test_gemini_image_to_video_kwargs_mapping(): with patch.object(nodes, "avideo_generation", gen_mock), patch.object( nodes, "avideo_content", content_mock - ), patch("asyncio.sleep", AsyncMock(return_value=None)): + ), patch.object( + nodes.GeminiModel, "image_tensor_to_bytesio", return_value=fake_bytesio + ) as tensor_mock, patch("asyncio.sleep", AsyncMock(return_value=None)): await nodes.APIVideoGenerate.execute(prompt="animate this photo", source=source) + tensor_mock.assert_called_once_with(fake_tensor) kwargs = gen_mock.call_args.kwargs assert kwargs["model"] == "gemini/veo-3.1-fast-generate-preview" assert kwargs["aspectRatio"] == "9:16" assert kwargs["resolution"] == "1080p" assert kwargs["durationSeconds"] == 8 assert kwargs["personGeneration"] == "allow_adult" - assert kwargs["input_reference"].endswith("my_photo.png") + assert kwargs["input_reference"] is fake_bytesio + + asyncio.run(run()) + + +def test_gemini_image_to_video_requires_reference_image(): + async def run(): + gen_mock = AsyncMock(return_value=FakeVideoObject("completed")) + + source = { + "source": "image", + "provider": { + "provider": "gemini", + "model": { + "model": "veo-3.1-fast-generate-preview", + "size": "9:16", + "resolution": "2K", + "seconds": 8, + "person_generation": "allow_adult", + }, + }, + } + + with patch.object(nodes, "avideo_generation", gen_mock): + with pytest.raises(ValueError, match="reference_image"): + await nodes.APIVideoGenerate.execute(prompt="animate this photo", source=source) + + gen_mock.assert_not_awaited() asyncio.run(run())