Update video node to support more parameters and make it more dynamic

This commit is contained in:
Andy Ratsirarson
2026-07-15 21:12:36 -07:00
parent 93dd43cd84
commit 05c29133de
11 changed files with 534 additions and 21 deletions
+5
View File
@@ -0,0 +1,5 @@
from src.models.openai import _SoraModel
class AzureModel(_SoraModel):
PROVIDER = "azure"
+65
View File
@@ -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.")
+84
View File
@@ -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
+30
View File
@@ -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"
+49
View File
@@ -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
View File
@@ -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)
+57
View File
@@ -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.",
)
+36
View File
@@ -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}).",
)
+43
View File
@@ -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).",
)
+26
View File
@@ -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
View File
@@ -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())