Allow input image socket node

This commit is contained in:
Andy Ratsirarson
2026-07-15 22:42:27 -07:00
parent 34b748fadc
commit 6691b570a7
6 changed files with 77 additions and 66 deletions
+1 -1
View File
@@ -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`.
+22 -3
View File
@@ -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
+4 -4
View File
@@ -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
+6 -14
View File
@@ -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.",
)
+8 -41
View File
@@ -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()
+36 -3
View File
@@ -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())