Update video node to support more parameters and make it more dynamic
This commit is contained in:
@@ -0,0 +1,5 @@
|
||||
from src.models.openai import _SoraModel
|
||||
|
||||
|
||||
class AzureModel(_SoraModel):
|
||||
PROVIDER = "azure"
|
||||
@@ -0,0 +1,65 @@
|
||||
from comfy_api.latest import io
|
||||
|
||||
|
||||
class Model:
|
||||
"""Base class for all provider/model integrations.
|
||||
|
||||
Subclasses set PROVIDER and VIDEO_SETTINGS (a settings.video.VideoSettings
|
||||
instance declaring which SIZE/RESOLUTIONS subset and duration range this
|
||||
provider supports), and implement list_models() + video_kwargs().
|
||||
"""
|
||||
|
||||
PROVIDER: str = None
|
||||
VIDEO_SETTINGS = None
|
||||
|
||||
@classmethod
|
||||
def list_models(cls) -> list[str]:
|
||||
"""Return this provider's available model ids for text-to-video."""
|
||||
raise NotImplementedError(f"list_models not implemented for {cls.__name__}.")
|
||||
|
||||
@classmethod
|
||||
def video_settings_inputs(cls):
|
||||
"""Build the io.Input list revealed once a specific model is selected."""
|
||||
if cls.VIDEO_SETTINGS is None:
|
||||
raise NotImplementedError(f"{cls.__name__} does not support video generation.")
|
||||
return [
|
||||
cls.VIDEO_SETTINGS.size_input(),
|
||||
cls.VIDEO_SETTINGS.resolution_input(),
|
||||
cls.VIDEO_SETTINGS.seconds_input(),
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def video_kwargs(cls, model_id: str, settings: dict) -> dict:
|
||||
"""Translate the generic {size, resolution, seconds, ...} dict captured by
|
||||
the node's DynamicCombo into the litellm avideo_generation kwargs this
|
||||
provider's API expects (e.g. Gemini's aspectRatio/resolution/durationSeconds)."""
|
||||
raise NotImplementedError(f"video_kwargs not implemented for {cls.__name__}.")
|
||||
|
||||
@classmethod
|
||||
def model_combo_option(cls):
|
||||
"""Build this provider's DynamicCombo.Option: a nested 'model' combo where
|
||||
each model reveals this provider's video settings."""
|
||||
return io.DynamicCombo.Option(
|
||||
cls.PROVIDER,
|
||||
[
|
||||
io.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
io.DynamicCombo.Option(model_id, cls.video_settings_inputs())
|
||||
for model_id in cls.list_models()
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
def image_to_image(self, *args, **kwargs):
|
||||
raise NotImplementedError("image_to_image not implemented for this model.")
|
||||
|
||||
def text_to_image(self, *args, **kwargs):
|
||||
raise NotImplementedError("text_to_image not implemented for this model.")
|
||||
|
||||
def text_to_video(self, *args, **kwargs):
|
||||
raise NotImplementedError("text_to_video not implemented for this model.")
|
||||
|
||||
def text_to_text(self, *args, **kwargs):
|
||||
raise NotImplementedError("text_to_text not implemented for this model.")
|
||||
@@ -0,0 +1,84 @@
|
||||
from comfy_api.latest import io
|
||||
|
||||
from src.models.base import Model
|
||||
from src.settings.base import Resolutions, Sizes
|
||||
from src.settings.video import VideoSettings
|
||||
|
||||
STATIC_MODELS = [
|
||||
"veo-3.1-fast-generate-preview",
|
||||
"veo-3.1-fast-generate-001",
|
||||
"veo-3.1-generate-preview",
|
||||
"veo-3.1-generate-001",
|
||||
"veo-3.1-lite-generate-preview",
|
||||
"veo-2.0-generate-001",
|
||||
]
|
||||
|
||||
RESOLUTION_MAP = {Resolutions._1K: "720p", Resolutions._2K: "1080p"}
|
||||
|
||||
|
||||
class GeminiModel(Model):
|
||||
"""Google Veo text-to-video via litellm's 'gemini/<model_id>' provider."""
|
||||
|
||||
PROVIDER = "gemini"
|
||||
VIDEO_SETTINGS = VideoSettings(
|
||||
SIZE=[Sizes._16_9, Sizes._9_16],
|
||||
RESOLUTIONS=[Resolutions._1K, Resolutions._2K],
|
||||
MIN_SECONDS=4,
|
||||
MAX_SECONDS=8,
|
||||
)
|
||||
|
||||
_models_cache = None
|
||||
|
||||
@classmethod
|
||||
def list_models(cls) -> list[str]:
|
||||
if cls._models_cache is not None:
|
||||
return cls._models_cache
|
||||
|
||||
try:
|
||||
from google import genai
|
||||
|
||||
models = []
|
||||
for model in genai.Client().models.list():
|
||||
name = getattr(model, "name", "") or ""
|
||||
if name.startswith("models/"):
|
||||
name = name[len("models/"):]
|
||||
if "veo" in name.lower():
|
||||
models.append(name)
|
||||
cls._models_cache = models or list(STATIC_MODELS)
|
||||
except Exception:
|
||||
cls._models_cache = list(STATIC_MODELS)
|
||||
|
||||
return cls._models_cache
|
||||
|
||||
@classmethod
|
||||
def video_settings_inputs(cls):
|
||||
inputs = super().video_settings_inputs()
|
||||
inputs.append(
|
||||
io.Combo.Input(
|
||||
"person_generation",
|
||||
options=["allow_all", "allow_adult"],
|
||||
default="allow_all",
|
||||
tooltip="Whether people may be generated in the video.",
|
||||
)
|
||||
)
|
||||
inputs.append(
|
||||
io.String.Input(
|
||||
"negative_prompt",
|
||||
default="",
|
||||
optional=True,
|
||||
tooltip="Things to avoid in the generated video.",
|
||||
)
|
||||
)
|
||||
return inputs
|
||||
|
||||
@classmethod
|
||||
def video_kwargs(cls, model_id: str, settings: dict) -> dict:
|
||||
kwargs = {
|
||||
"aspectRatio": settings["size"],
|
||||
"resolution": RESOLUTION_MAP[settings["resolution"]],
|
||||
"durationSeconds": settings["seconds"],
|
||||
"personGeneration": settings["person_generation"],
|
||||
}
|
||||
if settings.get("negative_prompt"):
|
||||
kwargs["negativePrompt"] = settings["negative_prompt"]
|
||||
return kwargs
|
||||
@@ -0,0 +1,30 @@
|
||||
from src.models.base import Model
|
||||
from src.settings.base import Resolutions, Sizes
|
||||
from src.settings.video import VideoSettings, dimensions_string
|
||||
|
||||
|
||||
class _SoraModel(Model):
|
||||
"""Shared Sora video logic for the OpenAI and Azure providers, which expose
|
||||
the same model family and parameters via litellm. Subclasses set PROVIDER."""
|
||||
|
||||
# MAX_SECONDS=12 is a reasonable-but-unconfirmed ceiling: litellm's own
|
||||
# default is "4" and no hard max was confirmed against Sora's API docs.
|
||||
VIDEO_SETTINGS = VideoSettings(
|
||||
SIZE=[Sizes._16_9, Sizes._9_16],
|
||||
RESOLUTIONS=[Resolutions._1K, Resolutions._2K],
|
||||
MIN_SECONDS=4,
|
||||
MAX_SECONDS=12,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def list_models(cls) -> list[str]:
|
||||
return ["sora-2", "sora-2-pro", "sora-2-pro-high-res"]
|
||||
|
||||
@classmethod
|
||||
def video_kwargs(cls, model_id: str, settings: dict) -> dict:
|
||||
size = dimensions_string(settings["size"], settings["resolution"], "x")
|
||||
return {"size": size, "seconds": str(settings["seconds"])}
|
||||
|
||||
|
||||
class OpenAIModel(_SoraModel):
|
||||
PROVIDER = "openai"
|
||||
@@ -0,0 +1,49 @@
|
||||
from comfy_api.latest import io
|
||||
|
||||
from src.models.base import Model
|
||||
from src.settings.base import Resolutions, Sizes
|
||||
from src.settings.video import VideoSettings, dimensions_string
|
||||
|
||||
|
||||
class RunwayMLModel(Model):
|
||||
PROVIDER = "runwayml"
|
||||
|
||||
VIDEO_SETTINGS = VideoSettings(
|
||||
SIZE=[Sizes._16_9, Sizes._9_16],
|
||||
RESOLUTIONS=[Resolutions._1K, Resolutions._2K],
|
||||
MIN_SECONDS=5,
|
||||
MAX_SECONDS=10,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def list_models(cls) -> list[str]:
|
||||
return ["gen3a_turbo", "gen4_turbo", "gen4_aleph", "gen4_image", "gen4_image_turbo"]
|
||||
|
||||
@classmethod
|
||||
def video_settings_inputs(cls):
|
||||
return [
|
||||
cls.VIDEO_SETTINGS.size_input(),
|
||||
cls.VIDEO_SETTINGS.resolution_input(),
|
||||
io.Combo.Input(
|
||||
"seconds",
|
||||
options=[5, 10],
|
||||
default=5,
|
||||
tooltip="Duration in seconds — RunwayML's turbo models support 5 or 10s.",
|
||||
),
|
||||
io.Int.Input(
|
||||
"seed",
|
||||
default=-1,
|
||||
min=-1,
|
||||
optional=True,
|
||||
tooltip="Random seed. -1 leaves it unset.",
|
||||
),
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def video_kwargs(cls, model_id: str, settings: dict) -> dict:
|
||||
ratio = dimensions_string(settings["size"], settings["resolution"], ":")
|
||||
return {
|
||||
"ratio": ratio,
|
||||
"duration": settings["seconds"],
|
||||
**({"seed": settings["seed"]} if settings.get("seed", -1) >= 0 else {}),
|
||||
}
|
||||
+40
-17
@@ -4,14 +4,18 @@ from io import BytesIO
|
||||
from comfy_api.latest import ComfyAPI, InputImpl, io
|
||||
from litellm import avideo_content, avideo_generation, avideo_status
|
||||
|
||||
from src.models.azure import AzureModel
|
||||
from src.models.gemini import GeminiModel
|
||||
from src.models.openai import OpenAIModel
|
||||
from src.models.runwayml import RunwayMLModel
|
||||
|
||||
api = ComfyAPI()
|
||||
|
||||
# litellm 1.92.0 implements video generation for these providers only.
|
||||
# Update this note (and re-check litellm's llms/<provider>/videos/ dirs) as support grows.
|
||||
SUPPORTED_PROVIDERS = "gemini, vertex_ai, runwayml, azure, openai"
|
||||
|
||||
_TERMINAL_STATUSES = ("completed", "failed")
|
||||
|
||||
_PROVIDERS = [GeminiModel, OpenAIModel, AzureModel, RunwayMLModel]
|
||||
_PROVIDERS_BY_NAME = {p.PROVIDER: p for p in _PROVIDERS}
|
||||
|
||||
|
||||
class APIVideoGenerate(io.ComfyNode):
|
||||
"""Text-to-video generation via litellm, using the caller's own provider API key (BYOK)."""
|
||||
@@ -24,24 +28,31 @@ class APIVideoGenerate(io.ComfyNode):
|
||||
category="external_api/video",
|
||||
description=(
|
||||
"Generates a video from a text prompt via litellm, using your own "
|
||||
f"provider API key. Providers litellm currently supports for video "
|
||||
f"generation: {SUPPORTED_PROVIDERS}."
|
||||
"provider API key. Pick a provider, then a model, to reveal its "
|
||||
"parameters."
|
||||
),
|
||||
inputs=[
|
||||
io.String.Input(
|
||||
"model",
|
||||
default="gemini/veo-3.1-fast-generate-preview",
|
||||
tooltip=(
|
||||
"litellm 'provider/model' string, e.g. "
|
||||
"gemini/veo-3.1-fast-generate-preview. Supported providers "
|
||||
f"today: {SUPPORTED_PROVIDERS}."
|
||||
),
|
||||
),
|
||||
io.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
tooltip="Text description of the video to generate.",
|
||||
),
|
||||
io.DynamicCombo.Input(
|
||||
"source",
|
||||
options=[
|
||||
io.DynamicCombo.Option(
|
||||
"text",
|
||||
[
|
||||
io.DynamicCombo.Input(
|
||||
"provider",
|
||||
options=[p.model_combo_option() for p in _PROVIDERS],
|
||||
tooltip="Video generation provider.",
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Generation source. Only text-to-video is supported today.",
|
||||
),
|
||||
io.String.Input(
|
||||
"api_key",
|
||||
default="",
|
||||
@@ -76,11 +87,23 @@ class APIVideoGenerate(io.ComfyNode):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(cls, model, prompt, api_key="", poll_interval=10.0, timeout=600.0):
|
||||
async def execute(cls, prompt, source, api_key="", poll_interval=10.0, timeout=600.0):
|
||||
node_id = cls.hidden.unique_id
|
||||
auth_kwargs = {"api_key": api_key} if api_key else {}
|
||||
|
||||
op = await avideo_generation(model=model, prompt=prompt, timeout=timeout, **auth_kwargs)
|
||||
provider_dict = source["provider"]
|
||||
provider_name = provider_dict["provider"]
|
||||
model_dict = provider_dict["model"]
|
||||
model_id = model_dict["model"]
|
||||
settings = {k: v for k, v in model_dict.items() if k != "model"}
|
||||
|
||||
provider_cls = _PROVIDERS_BY_NAME[provider_name]
|
||||
extra_kwargs = provider_cls.video_kwargs(model_id, settings)
|
||||
model = f"{provider_name}/{model_id}"
|
||||
|
||||
op = await avideo_generation(
|
||||
model=model, prompt=prompt, timeout=timeout, **auth_kwargs, **extra_kwargs
|
||||
)
|
||||
|
||||
elapsed = 0.0
|
||||
await api.execution.set_progress(0.0, timeout, node_id=node_id)
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
from enum import Enum
|
||||
|
||||
from comfy_api.latest import io
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class Sizes(str, Enum):
|
||||
"""Every aspect ratio used anywhere in the app. Add a new ratio here once;
|
||||
every model's SIZE list references these members instead of a raw string."""
|
||||
|
||||
_21_9 = "21:9"
|
||||
_16_9 = "16:9"
|
||||
_4_3 = "4:3"
|
||||
_1_1 = "1:1"
|
||||
_3_4 = "3:4"
|
||||
_9_16 = "9:16"
|
||||
_3_2 = "3:2"
|
||||
_5_4 = "5:4"
|
||||
_4_5 = "4:5"
|
||||
_2_3 = "2:3"
|
||||
|
||||
|
||||
class Resolutions(str, Enum):
|
||||
"""Every resolution tier used anywhere in the app. Add a new tier here
|
||||
once; every model's RESOLUTIONS list references these members instead of
|
||||
a raw string."""
|
||||
|
||||
_480 = "480"
|
||||
_720 = "720"
|
||||
_1080 = "1080"
|
||||
_1K = "1K"
|
||||
_2K = "2K"
|
||||
_4K = "4K"
|
||||
|
||||
|
||||
class Settings(BaseModel):
|
||||
"""Base config for a modality (image/video): the generic size/resolution
|
||||
options a model supports, plus the io.Input builders shared by both."""
|
||||
|
||||
SIZE: list[Sizes]
|
||||
RESOLUTIONS: list[Resolutions]
|
||||
|
||||
def size_input(self, default=None):
|
||||
return io.Combo.Input(
|
||||
"size",
|
||||
options=list(self.SIZE),
|
||||
default=default if default is not None else self.SIZE[0],
|
||||
tooltip="Aspect ratio of the generated output.",
|
||||
)
|
||||
|
||||
def resolution_input(self, default=None):
|
||||
return io.Combo.Input(
|
||||
"resolution",
|
||||
options=list(self.RESOLUTIONS),
|
||||
default=default if default is not None else self.RESOLUTIONS[0],
|
||||
tooltip="Output resolution tier.",
|
||||
)
|
||||
@@ -0,0 +1,36 @@
|
||||
from comfy_api.latest import io
|
||||
|
||||
from src.settings.base import Resolutions, Settings, Sizes
|
||||
|
||||
IMAGE_SIZES = [
|
||||
Sizes._21_9,
|
||||
Sizes._16_9,
|
||||
Sizes._3_2,
|
||||
Sizes._4_3,
|
||||
Sizes._5_4,
|
||||
Sizes._1_1,
|
||||
Sizes._4_5,
|
||||
Sizes._3_4,
|
||||
Sizes._2_3,
|
||||
Sizes._9_16,
|
||||
]
|
||||
IMAGE_RESOLUTIONS = [Resolutions._480, Resolutions._720, Resolutions._1080, Resolutions._4K]
|
||||
|
||||
DEFAULT_MIN_IMAGES = 1
|
||||
|
||||
|
||||
class ImageSettings(Settings):
|
||||
"""A model's image generation bounds: which of the master SIZE/RESOLUTIONS
|
||||
options it supports, and how many images it can generate per call."""
|
||||
|
||||
MIN_IMAGES: int = DEFAULT_MIN_IMAGES
|
||||
MAX_IMAGES: int
|
||||
|
||||
def num_images_input(self, default=None):
|
||||
return io.Int.Input(
|
||||
"num_images",
|
||||
default=default if default is not None else self.MIN_IMAGES,
|
||||
min=self.MIN_IMAGES,
|
||||
max=self.MAX_IMAGES,
|
||||
tooltip=f"Number of images to generate at once ({self.MIN_IMAGES}-{self.MAX_IMAGES}).",
|
||||
)
|
||||
@@ -0,0 +1,43 @@
|
||||
from comfy_api.latest import io
|
||||
|
||||
from src.settings.base import Resolutions, Settings, Sizes
|
||||
|
||||
VIDEO_SIZES = [Sizes._21_9, Sizes._16_9, Sizes._4_3, Sizes._1_1, Sizes._3_4, Sizes._9_16]
|
||||
VIDEO_RESOLUTIONS = [Resolutions._1K, Resolutions._2K, Resolutions._4K]
|
||||
|
||||
DEFAULT_MIN_SECONDS = 5
|
||||
|
||||
# Canonical (size, resolution) -> (width, height) pixel dimensions. Providers
|
||||
# that take a single literal dimension string (e.g. Sora's "1280x720", Runway's
|
||||
# "1280:720") format from this instead of keeping their own duplicate table.
|
||||
DIMENSIONS = {
|
||||
(Sizes._16_9, Resolutions._1K): (1280, 720),
|
||||
(Sizes._16_9, Resolutions._2K): (1920, 1080),
|
||||
(Sizes._9_16, Resolutions._1K): (720, 1280),
|
||||
(Sizes._9_16, Resolutions._2K): (1080, 1920),
|
||||
}
|
||||
|
||||
|
||||
def dimensions_string(size: str, resolution: str, separator: str) -> str:
|
||||
"""Format the canonical pixel dimensions for (size, resolution) using the
|
||||
given separator, e.g. dimensions_string("9:16", "2K", ":") -> "1080:1920"."""
|
||||
width, height = DIMENSIONS[(size, resolution)]
|
||||
return f"{width}{separator}{height}"
|
||||
|
||||
|
||||
class VideoSettings(Settings):
|
||||
"""A model's video generation bounds: which of the master SIZE/RESOLUTIONS
|
||||
options it supports, and its supported duration range."""
|
||||
|
||||
MIN_SECONDS: int = DEFAULT_MIN_SECONDS
|
||||
MAX_SECONDS: int
|
||||
|
||||
def seconds_input(self, default=None, step=1):
|
||||
return io.Int.Input(
|
||||
"seconds",
|
||||
default=default if default is not None else self.MIN_SECONDS,
|
||||
min=self.MIN_SECONDS,
|
||||
max=self.MAX_SECONDS,
|
||||
step=step,
|
||||
tooltip=f"Video duration in seconds ({self.MIN_SECONDS}-{self.MAX_SECONDS}s for this model).",
|
||||
)
|
||||
@@ -39,6 +39,29 @@ def _install_fake_comfy_api():
|
||||
class Float(_InputFactory):
|
||||
pass
|
||||
|
||||
class Int(_InputFactory):
|
||||
pass
|
||||
|
||||
class Combo(_InputFactory):
|
||||
pass
|
||||
|
||||
class _DynamicComboOption:
|
||||
def __init__(self, key, inputs):
|
||||
self.key = key
|
||||
self.inputs = inputs
|
||||
|
||||
class _DynamicComboInputSpec(_InputSpec):
|
||||
def __init__(self, name, options=None, **kwargs):
|
||||
super().__init__(name, options=options, **kwargs)
|
||||
self.options = options or []
|
||||
|
||||
class DynamicCombo:
|
||||
Option = _DynamicComboOption
|
||||
|
||||
@staticmethod
|
||||
def Input(id, options=None, **kwargs):
|
||||
return _DynamicComboInputSpec(id, options=options, **kwargs)
|
||||
|
||||
class Video(_OutputFactory):
|
||||
pass
|
||||
|
||||
@@ -59,6 +82,9 @@ def _install_fake_comfy_api():
|
||||
io = types.SimpleNamespace(
|
||||
String=String,
|
||||
Float=Float,
|
||||
Int=Int,
|
||||
Combo=Combo,
|
||||
DynamicCombo=DynamicCombo,
|
||||
Video=Video,
|
||||
Hidden=Hidden,
|
||||
Schema=Schema,
|
||||
|
||||
+99
-4
@@ -14,6 +14,24 @@ class FakeVideoObject:
|
||||
self.error = error
|
||||
|
||||
|
||||
def _gemini_source(model="veo-3.1-fast-generate-preview", **overrides):
|
||||
settings = {
|
||||
"model": model,
|
||||
"size": "16:9",
|
||||
"resolution": "1K",
|
||||
"seconds": 8,
|
||||
"person_generation": "allow_all",
|
||||
}
|
||||
settings.update(overrides)
|
||||
return {"source": "text", "provider": {"provider": "gemini", "model": settings}}
|
||||
|
||||
|
||||
def _openai_source(model="sora-2", **overrides):
|
||||
settings = {"model": model, "size": "16:9", "resolution": "1K", "seconds": 4}
|
||||
settings.update(overrides)
|
||||
return {"source": "text", "provider": {"provider": "openai", "model": settings}}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _hidden():
|
||||
nodes.APIVideoGenerate.hidden = types.SimpleNamespace(unique_id="test-node")
|
||||
@@ -35,8 +53,8 @@ def test_generate_success_polls_until_completed():
|
||||
"asyncio.sleep", AsyncMock(return_value=None)
|
||||
):
|
||||
result = await nodes.APIVideoGenerate.execute(
|
||||
model="gemini/veo-3.1-fast-generate-preview",
|
||||
prompt="a cat riding a bike",
|
||||
source=_gemini_source(),
|
||||
)
|
||||
|
||||
assert isinstance(result, nodes.io.NodeOutput)
|
||||
@@ -45,6 +63,11 @@ def test_generate_success_polls_until_completed():
|
||||
assert video.file_obj.getvalue() == b"fake-mp4-bytes"
|
||||
assert status_mock.call_count == 2
|
||||
gen_mock.assert_awaited_once()
|
||||
assert gen_mock.call_args.kwargs["model"] == "gemini/veo-3.1-fast-generate-preview"
|
||||
assert gen_mock.call_args.kwargs["aspectRatio"] == "16:9"
|
||||
assert gen_mock.call_args.kwargs["resolution"] == "720p"
|
||||
assert gen_mock.call_args.kwargs["durationSeconds"] == 8
|
||||
assert gen_mock.call_args.kwargs["personGeneration"] == "allow_all"
|
||||
assert "api_key" not in gen_mock.call_args.kwargs
|
||||
assert "api_key" not in status_mock.call_args.kwargs
|
||||
assert "api_key" not in content_mock.call_args.kwargs
|
||||
@@ -64,11 +87,14 @@ def test_api_key_threaded_through_every_call():
|
||||
"asyncio.sleep", AsyncMock(return_value=None)
|
||||
):
|
||||
await nodes.APIVideoGenerate.execute(
|
||||
model="openai/sora-2",
|
||||
prompt="a dog surfing",
|
||||
source=_openai_source(),
|
||||
api_key="sk-test-123",
|
||||
)
|
||||
|
||||
assert gen_mock.call_args.kwargs["model"] == "openai/sora-2"
|
||||
assert gen_mock.call_args.kwargs["size"] == "1280x720"
|
||||
assert gen_mock.call_args.kwargs["seconds"] == "4"
|
||||
assert gen_mock.call_args.kwargs["api_key"] == "sk-test-123"
|
||||
assert content_mock.call_args.kwargs["api_key"] == "sk-test-123"
|
||||
status_mock.assert_not_awaited() # already completed, no polling needed
|
||||
@@ -91,8 +117,8 @@ def test_failed_generation_raises_runtime_error():
|
||||
):
|
||||
with pytest.raises(RuntimeError, match="content filtered"):
|
||||
await nodes.APIVideoGenerate.execute(
|
||||
model="gemini/veo-3.1-fast-generate-preview",
|
||||
prompt="anything",
|
||||
source=_gemini_source(),
|
||||
)
|
||||
|
||||
content_mock.assert_not_awaited()
|
||||
@@ -110,10 +136,79 @@ def test_timeout_raises_when_generation_never_completes():
|
||||
), patch("asyncio.sleep", AsyncMock(return_value=None)):
|
||||
with pytest.raises(TimeoutError):
|
||||
await nodes.APIVideoGenerate.execute(
|
||||
model="gemini/veo-3.1-fast-generate-preview",
|
||||
prompt="anything",
|
||||
source=_gemini_source(),
|
||||
poll_interval=5.0,
|
||||
timeout=12.0,
|
||||
)
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_runwayml_kwargs_mapping():
|
||||
async def run():
|
||||
gen_mock = AsyncMock(return_value=FakeVideoObject("completed"))
|
||||
content_mock = AsyncMock(return_value=b"bytes")
|
||||
|
||||
source = {
|
||||
"source": "text",
|
||||
"provider": {
|
||||
"provider": "runwayml",
|
||||
"model": {
|
||||
"model": "gen3a_turbo",
|
||||
"size": "9:16",
|
||||
"resolution": "2K",
|
||||
"seconds": 10,
|
||||
"seed": 42,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
with patch.object(nodes, "avideo_generation", gen_mock), patch.object(
|
||||
nodes, "avideo_content", content_mock
|
||||
), patch("asyncio.sleep", AsyncMock(return_value=None)):
|
||||
await nodes.APIVideoGenerate.execute(prompt="anything", source=source)
|
||||
|
||||
assert gen_mock.call_args.kwargs["model"] == "runwayml/gen3a_turbo"
|
||||
assert gen_mock.call_args.kwargs["ratio"] == "1080:1920"
|
||||
assert gen_mock.call_args.kwargs["duration"] == 10
|
||||
assert gen_mock.call_args.kwargs["seed"] == 42
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_schema_builds_nested_source_provider_model_combo():
|
||||
schema = nodes.APIVideoGenerate.define_schema()
|
||||
source_input = next(i for i in schema.inputs if i.name == "source")
|
||||
assert [o.key for o in source_input.options] == ["text"]
|
||||
|
||||
provider_input = source_input.options[0].inputs[0]
|
||||
assert provider_input.name == "provider"
|
||||
provider_keys = [o.key for o in provider_input.options]
|
||||
assert provider_keys == ["gemini", "openai", "azure", "runwayml"]
|
||||
|
||||
for option in provider_input.options:
|
||||
model_input = option.inputs[0]
|
||||
assert model_input.name == "model"
|
||||
assert len(model_input.options) > 0
|
||||
for model_option in model_input.options:
|
||||
field_names = {i.name for i in model_option.inputs}
|
||||
assert {"size", "resolution", "seconds"} <= field_names
|
||||
|
||||
|
||||
def test_azure_provider_string_uses_azure_prefix():
|
||||
async def run():
|
||||
gen_mock = AsyncMock(return_value=FakeVideoObject("completed"))
|
||||
content_mock = AsyncMock(return_value=b"bytes")
|
||||
|
||||
source = _openai_source(model="sora-2-pro")
|
||||
source["provider"]["provider"] = "azure"
|
||||
|
||||
with patch.object(nodes, "avideo_generation", gen_mock), patch.object(
|
||||
nodes, "avideo_content", content_mock
|
||||
), patch("asyncio.sleep", AsyncMock(return_value=None)):
|
||||
await nodes.APIVideoGenerate.execute(prompt="anything", source=source)
|
||||
|
||||
assert gen_mock.call_args.kwargs["model"] == "azure/sora-2-pro"
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
Reference in New Issue
Block a user