Add Comfy Cloud API nodes

Amp-Thread-ID: https://ampcode.com/threads/T-019fca67-53e6-76d6-8561-0bb1b07473d1
Co-authored-by: Amp <amp@ampcode.com>

comfy_api_nodes/apis/comfy_cloud.py

@@ -0,0 +1,31 @@
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+ComfyCloudWorkflow = Literal["text-to-image", "text-to-video", "image-to-video", "image-edit"]
+
+
+class ComfyCloudWorkflowInputs(BaseModel):
+    prompt: str = Field(...)
+    image_url: str | None = Field(None)
+
+
+class ComfyCloudGenerateRequest(BaseModel):
+    workflow: ComfyCloudWorkflow = Field(...)
+    inputs: ComfyCloudWorkflowInputs = Field(...)
+
+
+class ComfyCloudGenerateResponse(BaseModel):
+    task_id: str = Field(...)
+    status: str = Field(...)
+    polling_url: str = Field(...)
+    cancel_url: str = Field(...)
+
+
+class ComfyCloudStatusResponse(BaseModel):
+    task_id: str = Field(...)
+    status: str = Field(...)
+    progress: float | None = Field(None)
+    output_url: str | None = Field(None)
+    error: str | None = Field(None)

comfy_api_nodes/nodes_comfy_cloud.py

@@ -0,0 +1,150 @@
+from typing import ClassVar
+
+from typing_extensions import override
+
+from comfy_api.latest import IO, ComfyExtension, Input
+from comfy_api_nodes.apis.comfy_cloud import (
+    ComfyCloudGenerateRequest,
+    ComfyCloudGenerateResponse,
+    ComfyCloudStatusResponse,
+    ComfyCloudWorkflow,
+    ComfyCloudWorkflowInputs,
+)
+from comfy_api_nodes.util import (
+    ApiEndpoint,
+    download_url_to_image_tensor,
+    download_url_to_video_output,
+    get_number_of_images,
+    poll_op,
+    sync_op,
+    upload_image_to_comfyapi,
+    validate_string,
+)
+
+
+_GENERATE_ENDPOINT = ApiEndpoint(path="/proxy/comfy-cloud/workflow/generate", method="POST")
+
+
+class _ComfyCloudWorkflowNode(IO.ComfyNode):
+    workflow: ClassVar[ComfyCloudWorkflow]
+    node_id: ClassVar[str]
+    display_name: ClassVar[str]
+    category: ClassVar[str]
+    requires_image: ClassVar[bool]
+    returns_video: ClassVar[bool]
+
+    @classmethod
+    def define_schema(cls) -> IO.Schema:
+        inputs = [
+            IO.String.Input(
+                "prompt",
+                multiline=True,
+                default="",
+                tooltip="Describe the content to generate or the edit to apply.",
+            )
+        ]
+        if cls.requires_image:
+            inputs.append(IO.Image.Input("image"))
+
+        output = IO.Video.Output() if cls.returns_video else IO.Image.Output()
+        return IO.Schema(
+            node_id=cls.node_id,
+            display_name=cls.display_name,
+            category=cls.category,
+            inputs=inputs,
+            outputs=[output],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+        )
+
+    @classmethod
+    async def execute(cls, prompt: str, image: Input.Image | None = None) -> IO.NodeOutput:
+        validate_string(prompt, min_length=1)
+
+        image_url = None
+        if cls.requires_image:
+            if get_number_of_images(image) != 1:
+                raise ValueError("Exactly one input image is required.")
+            image_url = await upload_image_to_comfyapi(cls, image)
+
+        task = await sync_op(
+            cls,
+            _GENERATE_ENDPOINT,
+            response_model=ComfyCloudGenerateResponse,
+            data=ComfyCloudGenerateRequest(
+                workflow=cls.workflow,
+                inputs=ComfyCloudWorkflowInputs(prompt=prompt, image_url=image_url),
+            ),
+        )
+        result = await poll_op(
+            cls,
+            ApiEndpoint(path=task.polling_url),
+            response_model=ComfyCloudStatusResponse,
+            status_extractor=lambda response: response.status,
+            progress_extractor=lambda response: response.progress,
+            cancel_endpoint=ApiEndpoint(path=task.cancel_url, method="POST"),
+        )
+        if not result.output_url:
+            detail = f": {result.error}" if result.error else ""
+            raise RuntimeError(f"Comfy Cloud task {result.task_id} completed without an output URL{detail}")
+
+        if cls.returns_video:
+            output = await download_url_to_video_output(result.output_url, cls=cls)
+        else:
+            output = await download_url_to_image_tensor(result.output_url, cls=cls)
+        return IO.NodeOutput(output)
+
+
+class ComfyCloudTextToImageNode(_ComfyCloudWorkflowNode):
+    workflow = "text-to-image"
+    node_id = "ComfyCloudTextToImageNode"
+    display_name = "Comfy Cloud Text to Image"
+    category = "partner/image/Comfy Cloud"
+    requires_image = False
+    returns_video = False
+
+
+class ComfyCloudTextToVideoNode(_ComfyCloudWorkflowNode):
+    workflow = "text-to-video"
+    node_id = "ComfyCloudTextToVideoNode"
+    display_name = "Comfy Cloud Text to Video"
+    category = "partner/video/Comfy Cloud"
+    requires_image = False
+    returns_video = True
+
+
+class ComfyCloudImageToVideoNode(_ComfyCloudWorkflowNode):
+    workflow = "image-to-video"
+    node_id = "ComfyCloudImageToVideoNode"
+    display_name = "Comfy Cloud Image to Video"
+    category = "partner/video/Comfy Cloud"
+    requires_image = True
+    returns_video = True
+
+
+class ComfyCloudImageEditNode(_ComfyCloudWorkflowNode):
+    workflow = "image-edit"
+    node_id = "ComfyCloudImageEditNode"
+    display_name = "Comfy Cloud Image Edit"
+    category = "partner/image/Comfy Cloud"
+    requires_image = True
+    returns_video = False
+
+
+class ComfyCloudExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
+        return [
+            ComfyCloudTextToImageNode,
+            ComfyCloudTextToVideoNode,
+            ComfyCloudImageToVideoNode,
+            ComfyCloudImageEditNode,
+        ]
+
+
+async def comfy_entrypoint() -> ComfyCloudExtension:
+    return ComfyCloudExtension()

tests-unit/comfy_api_nodes_test/comfy_cloud_test.py

@@ -0,0 +1,106 @@
+import asyncio
+from unittest.mock import AsyncMock
+
+import pytest
+import torch
+
+from comfy.cli_args import args
+
+if not torch.cuda.is_available():
+    args.cpu = True
+
+from comfy_api_nodes.apis.comfy_cloud import (
+    ComfyCloudGenerateRequest,
+    ComfyCloudGenerateResponse,
+    ComfyCloudStatusResponse,
+    ComfyCloudWorkflowInputs,
+)
+from comfy_api_nodes import nodes_comfy_cloud
+
+
+@pytest.mark.parametrize(
+    ("node", "workflow", "returns_video", "requires_image"),
+    [
+        (nodes_comfy_cloud.ComfyCloudTextToImageNode, "text-to-image", False, False),
+        (nodes_comfy_cloud.ComfyCloudTextToVideoNode, "text-to-video", True, False),
+        (nodes_comfy_cloud.ComfyCloudImageToVideoNode, "image-to-video", True, True),
+        (nodes_comfy_cloud.ComfyCloudImageEditNode, "image-edit", False, True),
+    ],
+)
+def test_workflow_submission_polling_and_download(monkeypatch, node, workflow, returns_video, requires_image):
+    sync = AsyncMock(
+        return_value=ComfyCloudGenerateResponse(
+            task_id="task-1",
+            status="queued",
+            polling_url="/proxy/comfy-cloud/workflow/tasks/task-1",
+            cancel_url="/proxy/comfy-cloud/workflow/tasks/task-1/cancel",
+        )
+    )
+    poll = AsyncMock(
+        return_value=ComfyCloudStatusResponse(
+            task_id="task-1",
+            status="completed",
+            progress=100,
+            output_url="https://example.com/output",
+        )
+    )
+    upload = AsyncMock(return_value="https://example.com/input.png")
+    image_download = AsyncMock(return_value="image-output")
+    video_download = AsyncMock(return_value="video-output")
+    monkeypatch.setattr(nodes_comfy_cloud, "sync_op", sync)
+    monkeypatch.setattr(nodes_comfy_cloud, "poll_op", poll)
+    monkeypatch.setattr(nodes_comfy_cloud, "upload_image_to_comfyapi", upload)
+    monkeypatch.setattr(nodes_comfy_cloud, "download_url_to_image_tensor", image_download)
+    monkeypatch.setattr(nodes_comfy_cloud, "download_url_to_video_output", video_download)
+    monkeypatch.setattr(nodes_comfy_cloud, "get_number_of_images", lambda image: 1)
+
+    image = object() if requires_image else None
+    output = asyncio.run(node.execute("A tiny fennec fox", image))
+
+    endpoint = sync.call_args.args[1]
+    request = sync.call_args.kwargs["data"]
+    assert endpoint.path == "/proxy/comfy-cloud/workflow/generate"
+    assert endpoint.method == "POST"
+    assert request == ComfyCloudGenerateRequest(
+        workflow=workflow,
+        inputs=ComfyCloudWorkflowInputs(
+            prompt="A tiny fennec fox",
+            image_url="https://example.com/input.png" if requires_image else None,
+        ),
+    )
+    assert upload.await_count == int(requires_image)
+
+    poll_endpoint = poll.call_args.args[1]
+    cancel_endpoint = poll.call_args.kwargs["cancel_endpoint"]
+    assert poll_endpoint.path == "/proxy/comfy-cloud/workflow/tasks/task-1"
+    assert cancel_endpoint.path == "/proxy/comfy-cloud/workflow/tasks/task-1/cancel"
+    assert cancel_endpoint.method == "POST"
+    assert output[0] == ("video-output" if returns_video else "image-output")
+
+
+@pytest.mark.parametrize(
+    "node",
+    [nodes_comfy_cloud.ComfyCloudImageToVideoNode, nodes_comfy_cloud.ComfyCloudImageEditNode],
+)
+def test_image_workflows_reject_batches(monkeypatch, node):
+    upload = AsyncMock()
+    monkeypatch.setattr(nodes_comfy_cloud, "get_number_of_images", lambda image: 2)
+    monkeypatch.setattr(nodes_comfy_cloud, "upload_image_to_comfyapi", upload)
+
+    with pytest.raises(ValueError, match="Exactly one input image"):
+        asyncio.run(node.execute("Animate this", object()))
+    upload.assert_not_awaited()
+
+
+def test_contract_omits_optional_status_fields():
+    request = ComfyCloudGenerateRequest(
+        workflow="text-to-image",
+        inputs=ComfyCloudWorkflowInputs(prompt="A lighthouse"),
+    )
+    status = ComfyCloudStatusResponse(task_id="task-1", status="queued")
+
+    assert request.model_dump(exclude_none=True) == {
+        "workflow": "text-to-image",
+        "inputs": {"prompt": "A lighthouse"},
+    }
+    assert status.model_dump(exclude_none=True) == {"task_id": "task-1", "status": "queued"}