From cd12eb4d23fb08b7effc6b76017220174ca3f68a Mon Sep 17 00:00:00 2001 From: Hunter Senft-Grupp Date: Wed, 5 Aug 2026 05:12:24 +0000 Subject: [PATCH 1/2] Add Comfy Cloud image POC nodes Amp-Thread-ID: https://ampcode.com/threads/T-019fd04d-4b31-73a3-8120-83bbf67ee979 Co-authored-by: Amp --- comfy_api_nodes/apis/comfy_cloud.py | 23 +- comfy_api_nodes/nodes_comfy_cloud.py | 262 +++++++++++++++++- .../comfy_api_nodes_test/comfy_cloud_test.py | 137 +++++++++ 3 files changed, 416 insertions(+), 6 deletions(-) diff --git a/comfy_api_nodes/apis/comfy_cloud.py b/comfy_api_nodes/apis/comfy_cloud.py index 20b60e3fc..bc7bc417f 100644 --- a/comfy_api_nodes/apis/comfy_cloud.py +++ b/comfy_api_nodes/apis/comfy_cloud.py @@ -3,12 +3,31 @@ from typing import Literal from pydantic import BaseModel, Field -ComfyCloudWorkflow = Literal["text-to-image", "text-to-video", "image-to-video", "image-edit"] +ComfyCloudWorkflow = Literal[ + "text-to-image", + "text-to-video", + "image-to-video", + "image-edit", + "image.ideogram-4-design.v1", + "image.krea-2-creative-image.v1", + "image.mage-flow-image.v1", + "image.flux-2-reference-edit.v1", + "image.qwen-image-edit-2511.v1", + "image.seedvr2-image-upscale.v1", +] class ComfyCloudWorkflowInputs(BaseModel): - prompt: str = Field(...) + prompt: str | None = Field(None) image_url: str | None = Field(None) + instruction: str | None = Field(None) + prompt_enhance: bool | None = Field(None) + negative_prompt: str | None = Field(None) + aspect_ratio: str | None = Field(None) + guidance: float | None = Field(None) + quality_mode: str | None = Field(None) + seed: int | None = Field(None, ge=0, le=0xFFFFFFFFFFFFFFFF) + scale: str | None = Field(None) class ComfyCloudGenerateRequest(BaseModel): diff --git a/comfy_api_nodes/nodes_comfy_cloud.py b/comfy_api_nodes/nodes_comfy_cloud.py index c49a5f8c5..6cfbb5abb 100644 --- a/comfy_api_nodes/nodes_comfy_cloud.py +++ b/comfy_api_nodes/nodes_comfy_cloud.py @@ -67,17 +67,25 @@ class _ComfyCloudWorkflowNode(IO.ComfyNode): 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) + image_url = await cls._upload_image(image) + return await cls._run(ComfyCloudWorkflowInputs(prompt=prompt, image_url=image_url)) + + @classmethod + async def _upload_image(cls, image: Input.Image, total_pixels: int | None = 2048 * 2048) -> str: + if get_number_of_images(image) != 1: + raise ValueError("Exactly one input image is required.") + return await upload_image_to_comfyapi(cls, image, total_pixels=total_pixels) + + @classmethod + async def _run(cls, inputs: ComfyCloudWorkflowInputs) -> IO.NodeOutput: task = await sync_op( cls, _GENERATE_ENDPOINT, response_model=ComfyCloudGenerateResponse, data=ComfyCloudGenerateRequest( workflow=cls.workflow, - inputs=ComfyCloudWorkflowInputs(prompt=prompt, image_url=image_url), + inputs=inputs, ), ) result = await poll_op( @@ -135,6 +143,246 @@ class ComfyCloudImageEditNode(_ComfyCloudWorkflowNode): returns_video = False +_ASPECT_RATIOS = ["1:1", "4:5", "3:4", "2:3", "3:2", "4:3", "16:9", "9:16"] +_UINT64_MAX = 0xFFFFFFFFFFFFFFFF + + +def _prompt_input(name: str = "prompt") -> IO.String.Input: + return IO.String.Input(name, multiline=True, default="") + + +def _aspect_ratio_input() -> IO.Combo.Input: + return IO.Combo.Input("aspect_ratio", options=_ASPECT_RATIOS, default="1:1") + + +def _seed_input() -> IO.Int.Input: + return IO.Int.Input("seed", default=0, min=0, max=_UINT64_MAX, control_after_generate=True) + + +def _image_schema(node_id: str, display_name: str, inputs: list[IO.Input]) -> IO.Schema: + return IO.Schema( + node_id=node_id, + display_name=display_name, + category="partner/image/Comfy Cloud", + inputs=inputs, + outputs=[IO.Image.Output()], + hidden=[ + IO.Hidden.auth_token_comfy_org, + IO.Hidden.api_key_comfy_org, + IO.Hidden.unique_id, + ], + is_api_node=True, + ) + + +class ComfyCloudIdeogram4DesignNode(_ComfyCloudWorkflowNode): + workflow = "image.ideogram-4-design.v1" + node_id = "ComfyCloudIdeogram4DesignNode" + display_name = "Ideogram 4 Design" + category = "partner/image/Comfy Cloud" + requires_image = False + returns_video = False + + @classmethod + def define_schema(cls) -> IO.Schema: + return _image_schema( + cls.node_id, + cls.display_name, + [ + _prompt_input(), + _aspect_ratio_input(), + IO.Combo.Input( + "quality_mode", options=["quality", "balanced", "fast"], default="balanced" + ), + _seed_input(), + ], + ) + + @classmethod + async def execute( + cls, prompt: str, aspect_ratio: str = "1:1", quality_mode: str = "balanced", seed: int = 0 + ) -> IO.NodeOutput: + validate_string(prompt, min_length=1, max_length=4096) + return await cls._run( + ComfyCloudWorkflowInputs( + prompt=prompt, aspect_ratio=aspect_ratio, quality_mode=quality_mode, seed=seed + ) + ) + + +class ComfyCloudKrea2CreativeImageNode(_ComfyCloudWorkflowNode): + workflow = "image.krea-2-creative-image.v1" + node_id = "ComfyCloudKrea2CreativeImageNode" + display_name = "Krea 2 Creative Image" + category = "partner/image/Comfy Cloud" + requires_image = False + returns_video = False + + @classmethod + def define_schema(cls) -> IO.Schema: + return _image_schema( + cls.node_id, + cls.display_name, + [ + _prompt_input(), + IO.Boolean.Input("prompt_enhance", default=True), + _aspect_ratio_input(), + _seed_input(), + ], + ) + + @classmethod + async def execute( + cls, prompt: str, prompt_enhance: bool = True, aspect_ratio: str = "1:1", seed: int = 0 + ) -> IO.NodeOutput: + validate_string(prompt, min_length=1, max_length=4096) + return await cls._run( + ComfyCloudWorkflowInputs( + prompt=prompt, prompt_enhance=prompt_enhance, aspect_ratio=aspect_ratio, seed=seed + ) + ) + + +class ComfyCloudMageFlowImageNode(_ComfyCloudWorkflowNode): + workflow = "image.mage-flow-image.v1" + node_id = "ComfyCloudMageFlowImageNode" + display_name = "Mage-Flow Image" + category = "partner/image/Comfy Cloud" + requires_image = False + returns_video = False + + @classmethod + def define_schema(cls) -> IO.Schema: + return _image_schema( + cls.node_id, + cls.display_name, + [ + _prompt_input(), + IO.String.Input("negative_prompt", multiline=True, default=""), + _aspect_ratio_input(), + _seed_input(), + ], + ) + + @classmethod + async def execute( + cls, prompt: str, negative_prompt: str = "", aspect_ratio: str = "1:1", seed: int = 0 + ) -> IO.NodeOutput: + validate_string(prompt, min_length=1, max_length=4096) + validate_string(negative_prompt, min_length=0, max_length=2048, field_name="negative_prompt") + return await cls._run( + ComfyCloudWorkflowInputs( + prompt=prompt, negative_prompt=negative_prompt, aspect_ratio=aspect_ratio, seed=seed + ) + ) + + +class ComfyCloudFlux2ReferenceEditNode(_ComfyCloudWorkflowNode): + workflow = "image.flux-2-reference-edit.v1" + node_id = "ComfyCloudFlux2ReferenceEditNode" + display_name = "FLUX.2 Reference Edit" + category = "partner/image/Comfy Cloud" + requires_image = True + returns_video = False + + @classmethod + def define_schema(cls) -> IO.Schema: + return _image_schema( + cls.node_id, + cls.display_name, + [ + IO.Image.Input("image"), + _prompt_input("instruction"), + IO.Float.Input("guidance", default=4.0, min=1.0, max=10.0, step=0.1), + IO.Combo.Input("quality_mode", options=["quality", "fast"], default="quality"), + _seed_input(), + ], + ) + + @classmethod + async def execute( + cls, + image: Input.Image, + instruction: str, + guidance: float = 4.0, + quality_mode: str = "quality", + seed: int = 0, + ) -> IO.NodeOutput: + validate_string(instruction, min_length=1, max_length=4096, field_name="instruction") + return await cls._run( + ComfyCloudWorkflowInputs( + image_url=await cls._upload_image(image, total_pixels=None), + instruction=instruction, + guidance=guidance, + quality_mode=quality_mode, + seed=seed, + ) + ) + + +class ComfyCloudQwenImageEdit2511Node(_ComfyCloudWorkflowNode): + workflow = "image.qwen-image-edit-2511.v1" + node_id = "ComfyCloudQwenImageEdit2511Node" + display_name = "Qwen Image Edit 2511" + category = "partner/image/Comfy Cloud" + requires_image = True + returns_video = False + + @classmethod + def define_schema(cls) -> IO.Schema: + return _image_schema( + cls.node_id, + cls.display_name, + [ + IO.Image.Input("image"), + _prompt_input("instruction"), + IO.Combo.Input("quality_mode", options=["quality", "fast"], default="quality"), + _seed_input(), + ], + ) + + @classmethod + async def execute( + cls, + image: Input.Image, + instruction: str, + quality_mode: str = "quality", + seed: int = 0, + ) -> IO.NodeOutput: + validate_string(instruction, min_length=1, max_length=4096, field_name="instruction") + return await cls._run( + ComfyCloudWorkflowInputs( + image_url=await cls._upload_image(image, total_pixels=None), + instruction=instruction, + quality_mode=quality_mode, + seed=seed, + ) + ) + + +class ComfyCloudSeedVR2ImageUpscaleNode(_ComfyCloudWorkflowNode): + workflow = "image.seedvr2-image-upscale.v1" + node_id = "ComfyCloudSeedVR2ImageUpscaleNode" + display_name = "SeedVR2 Image Upscale" + category = "partner/image/Comfy Cloud" + requires_image = True + returns_video = False + + @classmethod + def define_schema(cls) -> IO.Schema: + return _image_schema( + cls.node_id, + cls.display_name, + [IO.Image.Input("image"), IO.Combo.Input("scale", options=["2x", "4x"], default="4x")], + ) + + @classmethod + async def execute(cls, image: Input.Image, scale: str = "4x") -> IO.NodeOutput: + return await cls._run( + ComfyCloudWorkflowInputs(image_url=await cls._upload_image(image, total_pixels=None), scale=scale) + ) + + class ComfyCloudExtension(ComfyExtension): @override async def get_node_list(self) -> list[type[IO.ComfyNode]]: @@ -143,6 +391,12 @@ class ComfyCloudExtension(ComfyExtension): ComfyCloudTextToVideoNode, ComfyCloudImageToVideoNode, ComfyCloudImageEditNode, + ComfyCloudIdeogram4DesignNode, + ComfyCloudKrea2CreativeImageNode, + ComfyCloudMageFlowImageNode, + ComfyCloudFlux2ReferenceEditNode, + ComfyCloudQwenImageEdit2511Node, + ComfyCloudSeedVR2ImageUpscaleNode, ] diff --git a/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py b/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py index 91c37643b..c7a9299a3 100644 --- a/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py +++ b/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py @@ -1,5 +1,6 @@ import asyncio from io import BytesIO +from typing import get_args from unittest.mock import AsyncMock, Mock import pytest @@ -14,6 +15,7 @@ from comfy_api_nodes.apis.comfy_cloud import ( ComfyCloudGenerateRequest, ComfyCloudGenerateResponse, ComfyCloudStatusResponse, + ComfyCloudWorkflow, ComfyCloudWorkflowInputs, ) from comfy_api_nodes import nodes_comfy_cloud @@ -108,6 +110,141 @@ def test_contract_omits_optional_status_fields(): assert status.model_dump(exclude_none=True) == {"task_id": "task-1", "status": "queued"} +IMAGE_POC_NODES = [ + ( + nodes_comfy_cloud.ComfyCloudIdeogram4DesignNode, + "image.ideogram-4-design.v1", + ["prompt", "aspect_ratio", "quality_mode", "seed"], + { + "prompt": "A geometric fox logo", + "aspect_ratio": "4:5", + "quality_mode": "fast", + "seed": 11, + }, + ), + ( + nodes_comfy_cloud.ComfyCloudKrea2CreativeImageNode, + "image.krea-2-creative-image.v1", + ["prompt", "prompt_enhance", "aspect_ratio", "seed"], + {"prompt": "A glass forest", "prompt_enhance": False, "aspect_ratio": "16:9", "seed": 12}, + ), + ( + nodes_comfy_cloud.ComfyCloudMageFlowImageNode, + "image.mage-flow-image.v1", + ["prompt", "negative_prompt", "aspect_ratio", "seed"], + {"prompt": "A moonlit lake", "negative_prompt": "fog", "aspect_ratio": "3:2", "seed": 13}, + ), + ( + nodes_comfy_cloud.ComfyCloudFlux2ReferenceEditNode, + "image.flux-2-reference-edit.v1", + ["image", "instruction", "guidance", "quality_mode", "seed"], + {"image": object(), "instruction": "Make it winter", "guidance": 5.5, "quality_mode": "fast", "seed": 14}, + ), + ( + nodes_comfy_cloud.ComfyCloudQwenImageEdit2511Node, + "image.qwen-image-edit-2511.v1", + ["image", "instruction", "quality_mode", "seed"], + {"image": object(), "instruction": "Remove the sign", "quality_mode": "fast", "seed": 15}, + ), + ( + nodes_comfy_cloud.ComfyCloudSeedVR2ImageUpscaleNode, + "image.seedvr2-image-upscale.v1", + ["image", "scale"], + {"image": object(), "scale": "2x"}, + ), +] + + +@pytest.mark.parametrize(("node", "workflow", "input_names", "arguments"), IMAGE_POC_NODES) +def test_image_poc_node_schema_and_request_mapping(monkeypatch, node, workflow, input_names, arguments): + sync = AsyncMock( + return_value=ComfyCloudGenerateResponse( + task_id="task-poc", + status="queued", + polling_url="/tasks/task-poc", + cancel_url="/tasks/task-poc/cancel", + ) + ) + poll = AsyncMock( + return_value=ComfyCloudStatusResponse( + task_id="task-poc", status="completed", output_url="/results/task-poc/image.png" + ) + ) + upload = AsyncMock(return_value="/uploads/input.png") + download = AsyncMock(return_value="image-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", download) + monkeypatch.setattr(nodes_comfy_cloud, "get_number_of_images", lambda image: 1) + + schema = node.define_schema() + assert schema.node_id == node.node_id + assert schema.display_name == node.display_name + assert schema.is_api_node is True + assert [input.id for input in schema.inputs] == input_names + assert len(schema.outputs) == 1 + assert schema.outputs[0].get_io_type() == "IMAGE" + + output = asyncio.run(node.execute(**arguments)) + request = sync.call_args.kwargs["data"] + expected_inputs = {key: value for key, value in arguments.items() if key != "image"} + if "image" in arguments: + expected_inputs["image_url"] = "/uploads/input.png" + + assert request.workflow == workflow + assert request.inputs.model_dump(exclude_none=True) == expected_inputs + assert "asset_id" not in request.model_dump_json() + assert upload.await_count == int("image" in arguments) + if "image" in arguments: + assert upload.call_args.kwargs == {"total_pixels": None} + download.assert_awaited_once_with("/results/task-poc/image.png", cls=node) + assert output[0] == "image-output" + + +def test_image_poc_schema_defaults_ranges_and_enums(): + schemas = { + node.workflow: {input.id: input for input in node.define_schema().inputs} + for node, _, _, _ in IMAGE_POC_NODES + } + aspect_ratios = ["1:1", "4:5", "3:4", "2:3", "3:2", "4:3", "16:9", "9:16"] + + for workflow in [ + "image.ideogram-4-design.v1", + "image.krea-2-creative-image.v1", + "image.mage-flow-image.v1", + ]: + assert schemas[workflow]["aspect_ratio"].options == aspect_ratios + assert schemas[workflow]["aspect_ratio"].default == "1:1" + seed = schemas[workflow]["seed"] + assert (seed.default, seed.min, seed.max) == (0, 0, 0xFFFFFFFFFFFFFFFF) + + assert schemas["image.ideogram-4-design.v1"]["quality_mode"].options == ["quality", "balanced", "fast"] + assert schemas["image.ideogram-4-design.v1"]["quality_mode"].default == "balanced" + assert schemas["image.krea-2-creative-image.v1"]["prompt_enhance"].default is True + assert schemas["image.mage-flow-image.v1"]["negative_prompt"].default == "" + + guidance = schemas["image.flux-2-reference-edit.v1"]["guidance"] + assert (guidance.default, guidance.min, guidance.max, guidance.step) == (4.0, 1.0, 10.0, 0.1) + for workflow in ["image.flux-2-reference-edit.v1", "image.qwen-image-edit-2511.v1"]: + assert schemas[workflow]["quality_mode"].options == ["quality", "fast"] + assert schemas[workflow]["quality_mode"].default == "quality" + seed = schemas[workflow]["seed"] + assert (seed.default, seed.min, seed.max) == (0, 0, 0xFFFFFFFFFFFFFFFF) + + scale = schemas["image.seedvr2-image-upscale.v1"]["scale"] + assert scale.options == ["2x", "4x"] + assert scale.default == "4x" + + +def test_image_poc_api_declarations_and_extension_registration(): + workflows = {workflow for _, workflow, _, _ in IMAGE_POC_NODES} + registered = set(asyncio.run(nodes_comfy_cloud.ComfyCloudExtension().get_node_list())) + + assert workflows <= set(get_args(ComfyCloudWorkflow)) + assert {node for node, _, _, _ in IMAGE_POC_NODES} <= registered + + def test_download_cloud_audio_url_to_audio_input(monkeypatch): node = nodes_comfy_cloud.ComfyCloudTextToImageNode downloaded = b"encoded audio" From 42233d148d26c83d51e0a4957b3cbaae756efdd0 Mon Sep 17 00:00:00 2001 From: Hunter Senft-Grupp Date: Wed, 5 Aug 2026 05:34:47 +0000 Subject: [PATCH 2/2] Fix Image POC asset payloads Amp-Thread-ID: https://ampcode.com/threads/T-019fd04c-a312-75c9-ac5f-a1de1a7d4c4e Co-authored-by: Amp --- comfy_api_nodes/apis/comfy_cloud.py | 6 +++++ comfy_api_nodes/nodes_comfy_cloud.py | 22 ++++++++++++++++--- .../comfy_api_nodes_test/comfy_cloud_test.py | 3 ++- 3 files changed, 27 insertions(+), 4 deletions(-) diff --git a/comfy_api_nodes/apis/comfy_cloud.py b/comfy_api_nodes/apis/comfy_cloud.py index bc7bc417f..615740879 100644 --- a/comfy_api_nodes/apis/comfy_cloud.py +++ b/comfy_api_nodes/apis/comfy_cloud.py @@ -20,6 +20,7 @@ ComfyCloudWorkflow = Literal[ class ComfyCloudWorkflowInputs(BaseModel): prompt: str | None = Field(None) image_url: str | None = Field(None) + assets: dict[str, "ComfyCloudAssetInput"] | None = Field(None) instruction: str | None = Field(None) prompt_enhance: bool | None = Field(None) negative_prompt: str | None = Field(None) @@ -30,6 +31,11 @@ class ComfyCloudWorkflowInputs(BaseModel): scale: str | None = Field(None) +class ComfyCloudAssetInput(BaseModel): + type: Literal["IMAGE", "VIDEO", "AUDIO"] = Field(...) + url: str = Field(...) + + class ComfyCloudGenerateRequest(BaseModel): workflow: ComfyCloudWorkflow = Field(...) inputs: ComfyCloudWorkflowInputs = Field(...) diff --git a/comfy_api_nodes/nodes_comfy_cloud.py b/comfy_api_nodes/nodes_comfy_cloud.py index 6cfbb5abb..e82c1d854 100644 --- a/comfy_api_nodes/nodes_comfy_cloud.py +++ b/comfy_api_nodes/nodes_comfy_cloud.py @@ -4,6 +4,7 @@ from typing_extensions import override from comfy_api.latest import IO, ComfyExtension, Input from comfy_api_nodes.apis.comfy_cloud import ( + ComfyCloudAssetInput, ComfyCloudGenerateRequest, ComfyCloudGenerateResponse, ComfyCloudStatusResponse, @@ -311,7 +312,11 @@ class ComfyCloudFlux2ReferenceEditNode(_ComfyCloudWorkflowNode): validate_string(instruction, min_length=1, max_length=4096, field_name="instruction") return await cls._run( ComfyCloudWorkflowInputs( - image_url=await cls._upload_image(image, total_pixels=None), + assets={ + "image": ComfyCloudAssetInput( + type="IMAGE", url=await cls._upload_image(image, total_pixels=None) + ) + }, instruction=instruction, guidance=guidance, quality_mode=quality_mode, @@ -352,7 +357,11 @@ class ComfyCloudQwenImageEdit2511Node(_ComfyCloudWorkflowNode): validate_string(instruction, min_length=1, max_length=4096, field_name="instruction") return await cls._run( ComfyCloudWorkflowInputs( - image_url=await cls._upload_image(image, total_pixels=None), + assets={ + "image": ComfyCloudAssetInput( + type="IMAGE", url=await cls._upload_image(image, total_pixels=None) + ) + }, instruction=instruction, quality_mode=quality_mode, seed=seed, @@ -379,7 +388,14 @@ class ComfyCloudSeedVR2ImageUpscaleNode(_ComfyCloudWorkflowNode): @classmethod async def execute(cls, image: Input.Image, scale: str = "4x") -> IO.NodeOutput: return await cls._run( - ComfyCloudWorkflowInputs(image_url=await cls._upload_image(image, total_pixels=None), scale=scale) + ComfyCloudWorkflowInputs( + assets={ + "image": ComfyCloudAssetInput( + type="IMAGE", url=await cls._upload_image(image, total_pixels=None) + ) + }, + scale=scale, + ) ) diff --git a/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py b/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py index c7a9299a3..63e957451 100644 --- a/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py +++ b/tests-unit/comfy_api_nodes_test/comfy_cloud_test.py @@ -190,11 +190,12 @@ def test_image_poc_node_schema_and_request_mapping(monkeypatch, node, workflow, request = sync.call_args.kwargs["data"] expected_inputs = {key: value for key, value in arguments.items() if key != "image"} if "image" in arguments: - expected_inputs["image_url"] = "/uploads/input.png" + expected_inputs["assets"] = {"image": {"type": "IMAGE", "url": "/uploads/input.png"}} assert request.workflow == workflow assert request.inputs.model_dump(exclude_none=True) == expected_inputs assert "asset_id" not in request.model_dump_json() + assert '"id"' not in request.model_dump_json() assert upload.await_count == int("image" in arguments) if "image" in arguments: assert upload.call_args.kwargs == {"total_pixels": None}