diff --git a/comfy_api_nodes/apis/comfy_cloud.py b/comfy_api_nodes/apis/comfy_cloud.py new file mode 100644 index 000000000..20b60e3fc --- /dev/null +++ b/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) diff --git a/comfy_api_nodes/nodes_comfy_cloud.py b/comfy_api_nodes/nodes_comfy_cloud.py new file mode 100644 index 000000000..c49a5f8c5 --- /dev/null +++ b/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() diff --git a/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py b/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py new file mode 100644 index 000000000..1b6650b42 --- /dev/null +++ b/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"}