Andy RatsirarsonandClaude Sonnet 5
c24e649d1c
- New APIImageGenerate node (external_api/image): text-to-image and image-to-image via Gemini's Imagen (text-only) and Nano Banana / Nano Banana Pro / 2 (text and multi-reference image-to-image) models. - Omni Flash video now accepts a batch of reference images (not just one), sent in order as separate references in the same call. - Fix Veo/Omni resolution and duration handling: correct 720/1080/4K resolution options per model tier, drop the non-existent veo-3.1-fast-generate-001 id, note Google's 8s lock on 1080/4K. - Fix RunwayML model list: gen4_image/gen4_image_turbo are image models, not video, so drop them from the video node's provider list (tracked as follow-up work to add them to the image node). - Update README and tests accordingly. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014Bwg5aRmgChpwrVnpNgoQm
128 lines
4.5 KiB
Python
128 lines
4.5 KiB
Python
import asyncio
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
|
|
from src import nodes
|
|
|
|
|
|
class FakeImageObject:
|
|
def __init__(self, b64_json=None, url=None):
|
|
self.b64_json = b64_json
|
|
self.url = url
|
|
|
|
|
|
class FakeImageResponse:
|
|
def __init__(self, data):
|
|
self.data = data
|
|
|
|
|
|
def _text_source(model="imagen-4.0-generate-001", **overrides):
|
|
settings = {"model": model, "size": "1024x1024", "num_images": 1}
|
|
settings.update(overrides)
|
|
return {"source": "text", "provider": {"provider": "gemini", "model": settings}}
|
|
|
|
|
|
def _image_source(model="gemini-2.5-flash-image", **overrides):
|
|
settings = {"model": model, "size": "1024x1024", "num_images": 1}
|
|
settings.update(overrides)
|
|
return {"source": "image", "provider": {"provider": "gemini", "model": settings}}
|
|
|
|
|
|
def test_schema_lists_only_nano_banana_for_image_source():
|
|
schema = nodes.APIImageGenerate.define_schema()
|
|
source_input = next(i for i in schema.inputs if i.name == "source")
|
|
text_option, image_option = source_input.options
|
|
|
|
text_provider_input = text_option.inputs[0]
|
|
gemini_text_option = next(o for o in text_provider_input.options if o.key == "gemini")
|
|
text_model_ids = {o.key for o in gemini_text_option.inputs[0].options}
|
|
assert "imagen-4.0-generate-001" in text_model_ids
|
|
assert "gemini-2.5-flash-image" in text_model_ids
|
|
|
|
image_provider_input = image_option.inputs[0]
|
|
gemini_image_option = next(o for o in image_provider_input.options if o.key == "gemini")
|
|
image_model_ids = {o.key for o in gemini_image_option.inputs[0].options}
|
|
assert "gemini-2.5-flash-image" in image_model_ids
|
|
assert "imagen-4.0-generate-001" not in image_model_ids
|
|
|
|
for model_option in gemini_image_option.inputs[0].options:
|
|
field_names = {i.name for i in model_option.inputs}
|
|
assert {"size", "num_images", "reference_image"} <= field_names
|
|
|
|
|
|
def test_text_to_image_calls_aimage_generation():
|
|
async def run():
|
|
gen_mock = AsyncMock(
|
|
return_value=FakeImageResponse(data=[FakeImageObject(b64_json="ZmFrZS1wbmc=")])
|
|
)
|
|
|
|
with patch.object(nodes, "aimage_generation", gen_mock), patch.object(
|
|
nodes.Model,
|
|
"image_bytes_list_to_tensor",
|
|
return_value="fake-tensor",
|
|
) as tensor_mock:
|
|
result = await nodes.APIImageGenerate.execute(
|
|
prompt="a cat wearing a hat",
|
|
source=_text_source(),
|
|
)
|
|
|
|
gen_mock.assert_awaited_once()
|
|
assert gen_mock.call_args.kwargs["model"] == "gemini/imagen-4.0-generate-001"
|
|
assert gen_mock.call_args.kwargs["prompt"] == "a cat wearing a hat"
|
|
assert gen_mock.call_args.kwargs["size"] == "1024x1024"
|
|
assert gen_mock.call_args.kwargs["n"] == 1
|
|
tensor_mock.assert_called_once_with([b"fake-png"])
|
|
assert result.args[0] == "fake-tensor"
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_image_to_image_calls_aimage_edit_with_batched_images():
|
|
async def run():
|
|
edit_mock = AsyncMock(
|
|
return_value=FakeImageResponse(data=[FakeImageObject(b64_json="ZmFrZS1wbmc=")])
|
|
)
|
|
fake_tensor = object()
|
|
fake_buffers = [object(), object()]
|
|
|
|
source = _image_source()
|
|
source["provider"]["model"]["reference_image"] = fake_tensor
|
|
|
|
with patch.object(nodes, "aimage_edit", edit_mock), patch.object(
|
|
nodes.GeminiModel,
|
|
"image_tensor_to_bytesio_list",
|
|
return_value=fake_buffers,
|
|
) as split_mock, patch.object(
|
|
nodes.Model, "image_bytes_list_to_tensor", return_value="fake-tensor"
|
|
):
|
|
result = await nodes.APIImageGenerate.execute(
|
|
prompt="combine these two",
|
|
source=source,
|
|
)
|
|
|
|
split_mock.assert_called_once_with(fake_tensor)
|
|
edit_mock.assert_awaited_once()
|
|
assert edit_mock.call_args.kwargs["model"] == "gemini/gemini-2.5-flash-image"
|
|
assert edit_mock.call_args.kwargs["image"] == fake_buffers
|
|
assert edit_mock.call_args.kwargs["prompt"] == "combine these two"
|
|
assert result.args[0] == "fake-tensor"
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_image_to_image_requires_reference_image():
|
|
async def run():
|
|
edit_mock = AsyncMock()
|
|
|
|
with patch.object(nodes, "aimage_edit", edit_mock):
|
|
with pytest.raises(ValueError, match="reference_image"):
|
|
await nodes.APIImageGenerate.execute(
|
|
prompt="anything",
|
|
source=_image_source(),
|
|
)
|
|
|
|
edit_mock.assert_not_awaited()
|
|
|
|
asyncio.run(run())
|