Allow input image socket node
This commit is contained in:
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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())
|
||||
|
||||
|
||||
Reference in New Issue
Block a user