diff --git a/src/models/azure.py b/src/models/azure.py new file mode 100644 index 0000000..c21d2d1 --- /dev/null +++ b/src/models/azure.py @@ -0,0 +1,5 @@ +from src.models.openai import _SoraModel + + +class AzureModel(_SoraModel): + PROVIDER = "azure" diff --git a/src/models/base.py b/src/models/base.py new file mode 100644 index 0000000..8bd3643 --- /dev/null +++ b/src/models/base.py @@ -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.") \ No newline at end of file diff --git a/src/models/gemini.py b/src/models/gemini.py new file mode 100644 index 0000000..f97f759 --- /dev/null +++ b/src/models/gemini.py @@ -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/' 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 diff --git a/src/models/openai.py b/src/models/openai.py new file mode 100644 index 0000000..7907772 --- /dev/null +++ b/src/models/openai.py @@ -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" diff --git a/src/models/runwayml.py b/src/models/runwayml.py new file mode 100644 index 0000000..934a04d --- /dev/null +++ b/src/models/runwayml.py @@ -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 {}), + } diff --git a/src/nodes.py b/src/nodes.py index 5485048..4004281 100644 --- a/src/nodes.py +++ b/src/nodes.py @@ -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//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) diff --git a/src/settings/base.py b/src/settings/base.py new file mode 100644 index 0000000..d1ac3c7 --- /dev/null +++ b/src/settings/base.py @@ -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.", + ) diff --git a/src/settings/image.py b/src/settings/image.py new file mode 100644 index 0000000..d368c68 --- /dev/null +++ b/src/settings/image.py @@ -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}).", + ) diff --git a/src/settings/video.py b/src/settings/video.py new file mode 100644 index 0000000..0eb4e88 --- /dev/null +++ b/src/settings/video.py @@ -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).", + ) diff --git a/tests/conftest.py b/tests/conftest.py index 60b8378..2abbe79 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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, diff --git a/tests/test_nodes.py b/tests/test_nodes.py index 3f1eb9d..7372c8a 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -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())