zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-08-12 04:43:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge Image POC ComfyUI nodes

Browse Source
This commit is contained in:
Hunter Senft-Grupp
2026-08-05 05:48:53 +00:00
parent 93b47ed9fb 42233d148d
commit d5d1d0655c
3 changed files with 439 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
comfy_api_nodes/apis/comfy_cloud.py
View File

@@ -3,12 +3,37 @@ 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)
assets: dict[str, "ComfyCloudAssetInput"] | 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 ComfyCloudAssetInput(BaseModel):
type: Literal["IMAGE", "VIDEO", "AUDIO"] = Field(...)
url: str = Field(...)
class ComfyCloudGenerateRequest(BaseModel):

278
comfy_api_nodes/nodes_comfy_cloud.py
View File

@@ -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,
@@ -67,17 +68,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 +144,261 @@ 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(
assets={
"image": ComfyCloudAssetInput(
type="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(
assets={
"image": ComfyCloudAssetInput(
type="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(
assets={
"image": ComfyCloudAssetInput(
type="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 +407,12 @@ class ComfyCloudExtension(ComfyExtension):
ComfyCloudTextToVideoNode,
ComfyCloudImageToVideoNode,
ComfyCloudImageEditNode,
ComfyCloudIdeogram4DesignNode,
ComfyCloudKrea2CreativeImageNode,
ComfyCloudMageFlowImageNode,
ComfyCloudFlux2ReferenceEditNode,
ComfyCloudQwenImageEdit2511Node,
ComfyCloudSeedVR2ImageUpscaleNode,
]

138
tests-unit/comfy_api_nodes_test/comfy_cloud_test.py
View File

@@ -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,142 @@ 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["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}
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"