diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..75027e2 --- /dev/null +++ b/__init__.py @@ -0,0 +1,15 @@ +from typing_extensions import override + +from comfy_api.latest import ComfyExtension, io + +from .src.nodes import NODES + + +class ExternalApiExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return NODES + + +async def comfy_entrypoint() -> ExternalApiExtension: + return ExternalApiExtension() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..deb4ef8 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,44 @@ +[project] +name = "externalapi" +description = "ComfyUI node package for BYOK (bring-your-own-key) access to external generative APIs: Gemini, Veo, and more via LiteLLM." +version = "1.0.0" +license = "MIT" +# classifiers = [ +# # For OS-independent nodes (works on all operating systems) +# "Operating System :: OS Independent", +# +# # OR for OS-specific nodes, specify the supported systems: +# "Operating System :: Microsoft :: Windows", # Windows specific +# "Operating System :: POSIX :: Linux", # Linux specific +# "Operating System :: MacOS", # macOS specific +# +# # GPU Accelerator support. Pick the ones that are supported by your extension. +# "Environment :: GPU :: NVIDIA CUDA", # NVIDIA CUDA support +# "Environment :: GPU :: AMD ROCm", # AMD ROCm support +# "Environment :: GPU :: Intel Arc", # Intel Arc support +# "Environment :: NPU :: Huawei Ascend", # Huawei Ascend support +# "Environment :: GPU :: Apple Metal", # Apple Metal support +# ] + +dependencies = ["litellm==1.92.0", "google-genai==1.47.0"] + +[project.optional-dependencies] +dev = ["pytest>=7.0"] + +[tool.pytest.ini_options] +pythonpath = ["."] +testpaths = ["tests"] + +[project.urls] +Repository = "https://github.com/andy-ratsirarson/ComfyUI_ExternalAPI" +# Used by Comfy Registry https://registry.comfy.org +Documentation = "https://github.com/andy-ratsirarson/ComfyUI_ExternalAPI/wiki" +"Bug Tracker" = "https://github.com/andy-ratsirarson/ComfyUI_ExternalAPI/issues" + +[tool.comfy] +PublisherId = "" +DisplayName = "ComfyUI_ExternalAPI" +Icon = "" +includes = [] +# "requires-comfyui" = ">=1.0.0" # ComfyUI version compatibility + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..c0dec69 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +litellm==1.92.0 +google-genai==1.47.0 \ No newline at end of file diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/nodes.py b/src/nodes.py new file mode 100644 index 0000000..5485048 --- /dev/null +++ b/src/nodes.py @@ -0,0 +1,104 @@ +import asyncio +from io import BytesIO + +from comfy_api.latest import ComfyAPI, InputImpl, io +from litellm import avideo_content, avideo_generation, avideo_status + +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") + + +class APIVideoGenerate(io.ComfyNode): + """Text-to-video generation via litellm, using the caller's own provider API key (BYOK).""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="APIVideoGenerate", + display_name="API Video Generate (BYOK)", + 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}." + ), + 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.String.Input( + "api_key", + default="", + optional=True, + tooltip=( + "Overrides the provider's standard env var (e.g. GEMINI_API_KEY, " + "OPENAI_API_KEY) when set. Warning: unlike an env var, this value " + "is saved into the workflow JSON and into output file metadata. " + "Prefer leaving this blank and setting the env var on the ComfyUI " + "server instead." + ), + ), + io.Float.Input( + "poll_interval", + default=10.0, + min=1.0, + optional=True, + tooltip="Seconds between status checks while the video is generating.", + ), + io.Float.Input( + "timeout", + default=600.0, + min=10.0, + optional=True, + tooltip="Maximum seconds to wait for generation before giving up.", + ), + ], + outputs=[ + io.Video.Output(), + ], + hidden=[io.Hidden.unique_id], + ) + + @classmethod + async def execute(cls, model, prompt, 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) + + elapsed = 0.0 + await api.execution.set_progress(0.0, timeout, node_id=node_id) + while op.status not in _TERMINAL_STATUSES: + if elapsed >= timeout: + raise TimeoutError(f"Video generation timed out after {timeout}s") + await asyncio.sleep(poll_interval) + elapsed += poll_interval + await api.execution.set_progress(min(elapsed, timeout), timeout, node_id=node_id) + op = await avideo_status(video_id=op.id, timeout=timeout, **auth_kwargs) + + if op.status == "failed": + raise RuntimeError(f"Video generation failed: {op.error}") + + video_bytes = await avideo_content(video_id=op.id, **auth_kwargs) + await api.execution.set_progress(timeout, timeout, node_id=node_id) + + return io.NodeOutput(InputImpl.VideoFromFile(BytesIO(video_bytes))) + + +NODES = [APIVideoGenerate] diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..60b8378 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,101 @@ +"""Stubs the parts of `comfy_api.latest` our nodes touch. + +`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 sys +import types + + +def _install_fake_comfy_api(): + if "comfy_api.latest" in sys.modules: + return + + class _InputSpec: + def __init__(self, name, **kwargs): + self.name = name + self.kwargs = kwargs + + class _InputFactory: + @staticmethod + def Input(name, **kwargs): + return _InputSpec(name, **kwargs) + + class _OutputSpec: + def __init__(self, name=None, **kwargs): + self.name = name + self.kwargs = kwargs + + class _OutputFactory: + @staticmethod + def Output(name=None, **kwargs): + return _OutputSpec(name, **kwargs) + + class String(_InputFactory): + pass + + class Float(_InputFactory): + pass + + class Video(_OutputFactory): + pass + + class Hidden: + unique_id = "unique_id" + + class Schema: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class NodeOutput: + def __init__(self, *args): + self.args = args + + class ComfyNode: + hidden = None + + io = types.SimpleNamespace( + String=String, + Float=Float, + Video=Video, + Hidden=Hidden, + Schema=Schema, + NodeOutput=NodeOutput, + ComfyNode=ComfyNode, + ) + + class ComfyExtension: + async def get_node_list(self): + raise NotImplementedError + + class _Execution: + async def set_progress(self, value, max_value, node_id=None, **kwargs): + return None + + class ComfyAPI: + def __init__(self): + self.execution = _Execution() + + class VideoFromFile: + def __init__(self, file_obj): + self.file_obj = file_obj + + InputImpl = types.SimpleNamespace(VideoFromFile=VideoFromFile) + + latest = types.SimpleNamespace( + io=io, + ComfyExtension=ComfyExtension, + ComfyAPI=ComfyAPI, + InputImpl=InputImpl, + ) + + comfy_api = types.ModuleType("comfy_api") + comfy_api.latest = latest + + sys.modules["comfy_api"] = comfy_api + sys.modules["comfy_api.latest"] = latest + + +_install_fake_comfy_api() diff --git a/tests/test_nodes.py b/tests/test_nodes.py new file mode 100644 index 0000000..3f1eb9d --- /dev/null +++ b/tests/test_nodes.py @@ -0,0 +1,119 @@ +import asyncio +import types +from unittest.mock import AsyncMock, patch + +import pytest + +from src import nodes + + +class FakeVideoObject: + def __init__(self, status, id="vid_1", error=None): + self.status = status + self.id = id + self.error = error + + +@pytest.fixture(autouse=True) +def _hidden(): + nodes.APIVideoGenerate.hidden = types.SimpleNamespace(unique_id="test-node") + yield + nodes.APIVideoGenerate.hidden = None + + +def test_generate_success_polls_until_completed(): + async def run(): + gen_mock = AsyncMock(return_value=FakeVideoObject("processing")) + status_mock = AsyncMock( + side_effect=[FakeVideoObject("processing"), FakeVideoObject("completed")] + ) + content_mock = AsyncMock(return_value=b"fake-mp4-bytes") + + with patch.object(nodes, "avideo_generation", gen_mock), patch.object( + nodes, "avideo_status", status_mock + ), patch.object(nodes, "avideo_content", content_mock), patch( + "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", + ) + + assert isinstance(result, nodes.io.NodeOutput) + video = result.args[0] + assert isinstance(video, nodes.InputImpl.VideoFromFile) + assert video.file_obj.getvalue() == b"fake-mp4-bytes" + assert status_mock.call_count == 2 + gen_mock.assert_awaited_once() + 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 + + asyncio.run(run()) + + +def test_api_key_threaded_through_every_call(): + async def run(): + gen_mock = AsyncMock(return_value=FakeVideoObject("completed")) + status_mock = AsyncMock() + content_mock = AsyncMock(return_value=b"bytes") + + with patch.object(nodes, "avideo_generation", gen_mock), patch.object( + nodes, "avideo_status", status_mock + ), patch.object(nodes, "avideo_content", content_mock), patch( + "asyncio.sleep", AsyncMock(return_value=None) + ): + await nodes.APIVideoGenerate.execute( + model="openai/sora-2", + prompt="a dog surfing", + api_key="sk-test-123", + ) + + 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 + + asyncio.run(run()) + + +def test_failed_generation_raises_runtime_error(): + async def run(): + gen_mock = AsyncMock(return_value=FakeVideoObject("processing")) + status_mock = AsyncMock( + return_value=FakeVideoObject("failed", error={"message": "content filtered"}) + ) + content_mock = AsyncMock() + + with patch.object(nodes, "avideo_generation", gen_mock), patch.object( + nodes, "avideo_status", status_mock + ), patch.object(nodes, "avideo_content", content_mock), patch( + "asyncio.sleep", AsyncMock(return_value=None) + ): + with pytest.raises(RuntimeError, match="content filtered"): + await nodes.APIVideoGenerate.execute( + model="gemini/veo-3.1-fast-generate-preview", + prompt="anything", + ) + + content_mock.assert_not_awaited() + + asyncio.run(run()) + + +def test_timeout_raises_when_generation_never_completes(): + async def run(): + gen_mock = AsyncMock(return_value=FakeVideoObject("processing")) + status_mock = AsyncMock(return_value=FakeVideoObject("processing")) + + with patch.object(nodes, "avideo_generation", gen_mock), patch.object( + nodes, "avideo_status", status_mock + ), 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", + poll_interval=5.0, + timeout=12.0, + ) + + asyncio.run(run())