feat(api-nodes): add ByteDance Image nodes

comfy_api_nodes/nodes_bytedance.py

@@ -0,0 +1,336 @@
+import logging
+from enum import Enum
+from typing import Optional
+from typing_extensions import override
+
+import torch
+from pydantic import BaseModel, Field
+
+from comfy_api.latest import ComfyExtension, io as comfy_io
+from comfy_api_nodes.util.validation_utils import (
+    validate_image_aspect_ratio_range,
+    get_number_of_images,
+)
+from comfy_api_nodes.apis.client import (
+    ApiEndpoint,
+    HttpMethod,
+    SynchronousOperation,
+)
+from comfy_api_nodes.apinode_utils import download_url_to_image_tensor, upload_images_to_comfyapi, validate_string
+
+
+BYTEPLUS_ENDPOINT = "/proxy/byteplus/api/v3/images/generations"
+
+
+class Text2ImageModelName(str, Enum):
+    seedream3 = "seedream-3-0-t2i-250415"
+
+
+class Image2ImageModelName(str, Enum):
+    seededit3 = "seededit-3-0-i2i-250628"
+
+
+class Text2ImageTaskCreationRequest(BaseModel):
+    model: Text2ImageModelName = Text2ImageModelName.seedream3
+    prompt: str = Field(...)
+    response_format: Optional[str] = Field("url")
+    size: Optional[str] = Field(None)
+    seed: Optional[int] = Field(0, ge=0, le=2147483647)
+    guidance_scale: Optional[float] = Field(..., ge=1.0, le=10.0)
+    watermark: Optional[bool] = Field(True)
+
+
+class Image2ImageTaskCreationRequest(BaseModel):
+    model: Image2ImageModelName = Image2ImageModelName.seededit3
+    prompt: str = Field(...)
+    response_format: Optional[str] = Field("url")
+    image: str = Field(..., description="Base64 encoded string or image URL")
+    size: Optional[str] = Field("adaptive")
+    seed: Optional[int] = Field(..., ge=0, le=2147483647)
+    guidance_scale: Optional[float] = Field(..., ge=1.0, le=10.0)
+    watermark: Optional[bool] = Field(True)
+
+
+class ImageTaskCreationResponse(BaseModel):
+    model: str = Field(...)
+    created: int = Field(..., description="Unix timestamp (in seconds) indicating time when the request was created.")
+    data: list = Field([], description="Contains information about the generated image(s).")
+    error: dict = Field({}, description="Contains `code` and `message` fields in case of error.")
+
+
+RECOMMENDED_PRESETS = [
+    ("1024x1024 (1:1)", 1024, 1024),
+    ("864x1152 (3:4)", 864, 1152),
+    ("1152x864 (4:3)", 1152, 864),
+    ("1280x720 (16:9)", 1280, 720),
+    ("720x1280 (9:16)", 720, 1280),
+    ("832x1248 (2:3)", 832, 1248),
+    ("1248x832 (3:2)", 1248, 832),
+    ("1512x648 (21:9)", 1512, 648),
+    ("2048x2048 (1:1)", 2048, 2048),
+    ("Custom", None, None),
+]
+
+
+def get_image_url_from_response(response: ImageTaskCreationResponse) -> str:
+    if response.error:
+        error_msg = f"ByteDance request failed. Code: {response.error['code']}, message: {response.error['message']}"
+        logging.info(error_msg)
+        raise RuntimeError(error_msg)
+    logging.info("ByteDance task succeeded, image URL: %s", response.data[0]["url"])
+    return response.data[0]["url"]
+
+
+class ByteDanceImageNode(comfy_io.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return comfy_io.Schema(
+            node_id="ByteDanceImageNode",
+            display_name="ByteDance Image",
+            category="api node/image/ByteDance",
+            description="Generate images using ByteDance models via api based on prompt",
+            inputs=[
+                comfy_io.Combo.Input(
+                    "model",
+                    options=[model.value for model in Text2ImageModelName],
+                    default=Text2ImageModelName.seedream3.value,
+                    tooltip="Model name",
+                ),
+                comfy_io.String.Input(
+                    "prompt",
+                    multiline=True,
+                    tooltip="The text prompt used to generate the image",
+                ),
+                comfy_io.Combo.Input(
+                    "size_preset",
+                    options=[label for label, _, _ in RECOMMENDED_PRESETS],
+                    tooltip="Pick a recommended size. Select Custom to use the width and height below",
+                ),
+                comfy_io.Int.Input(
+                    "width",
+                    default=1024,
+                    min=512,
+                    max=2048,
+                    step=64,
+                    tooltip="Custom width for image. Value is working only if `size_preset` is set to `Custom`",
+                ),
+                comfy_io.Int.Input(
+                    "height",
+                    default=1024,
+                    min=512,
+                    max=2048,
+                    step=64,
+                    tooltip="Custom height for image. Value is working only if `size_preset` is set to `Custom`",
+                ),
+                comfy_io.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=2147483647,
+                    step=1,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    control_after_generate=True,
+                    tooltip="Seed to use for generation",
+                    optional=True,
+                ),
+                comfy_io.Float.Input(
+                    "guidance_scale",
+                    default=2.5,
+                    min=1.0,
+                    max=10.0,
+                    step=0.01,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    tooltip="Higher value makes the image follow the prompt more closely",
+                    optional=True,
+                ),
+                comfy_io.Boolean.Input(
+                    "watermark",
+                    default=True,
+                    tooltip="Whether to add an \"AI generated\" watermark to the image",
+                    optional=True,
+                ),
+            ],
+            outputs=[
+                comfy_io.Image.Output(),
+            ],
+            hidden=[
+                comfy_io.Hidden.auth_token_comfy_org,
+                comfy_io.Hidden.api_key_comfy_org,
+                comfy_io.Hidden.unique_id,
+            ],
+            is_api_node=True,
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        model: str,
+        prompt: str,
+        size_preset: str,
+        width: int,
+        height: int,
+        seed: int,
+        guidance_scale: float,
+        watermark: bool,
+    ) -> comfy_io.NodeOutput:
+        validate_string(prompt, strip_whitespace=True, min_length=1)
+        w = h = None
+        for label, tw, th in RECOMMENDED_PRESETS:
+            if label == size_preset:
+                w, h = tw, th
+                break
+
+        if w is None or h is None:
+            w, h = width, height
+            if not (512 <= w <= 2048) or not (512 <= h <= 2048):
+                raise ValueError(
+                    f"Custom size out of range: {w}x{h}. "
+                    "Both width and height must be between 512 and 2048 pixels."
+                )
+
+        payload = Text2ImageTaskCreationRequest(
+            model=model,
+            prompt=prompt,
+            size=f"{w}x{h}",
+            seed=seed,
+            guidance_scale=guidance_scale,
+            watermark=watermark,
+        )
+        auth_kwargs = {
+            "auth_token": cls.hidden.auth_token_comfy_org,
+            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        }
+        response = await SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path=BYTEPLUS_ENDPOINT,
+                method=HttpMethod.POST,
+                request_model=Text2ImageTaskCreationRequest,
+                response_model=ImageTaskCreationResponse,
+            ),
+            request=payload,
+            auth_kwargs=auth_kwargs,
+        ).execute()
+        return comfy_io.NodeOutput(await download_url_to_image_tensor(get_image_url_from_response(response)))
+
+
+class ByteDanceImageEditNode(comfy_io.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return comfy_io.Schema(
+            node_id="ByteDanceImageEditNode",
+            display_name="ByteDance Image Edit",
+            category="api node/video/ByteDance",
+            description="Edit images using ByteDance models via api based on prompt",
+            inputs=[
+                comfy_io.Combo.Input(
+                    "model",
+                    options=[model.value for model in Image2ImageModelName],
+                    default=Image2ImageModelName.seededit3.value,
+                    tooltip="Model name",
+                ),
+                comfy_io.Image.Input(
+                    "image",
+                    tooltip="The base image to edit",
+                ),
+                comfy_io.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="Instruction to edit image",
+                ),
+                comfy_io.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=2147483647,
+                    step=1,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    control_after_generate=True,
+                    tooltip="Seed to use for generation",
+                    optional=True,
+                ),
+                comfy_io.Float.Input(
+                    "guidance_scale",
+                    default=5.5,
+                    min=1.0,
+                    max=10.0,
+                    step=0.01,
+                    display_mode=comfy_io.NumberDisplay.number,
+                    tooltip="Higher value makes the image follow the prompt more closely",
+                    optional=True,
+                ),
+                comfy_io.Boolean.Input(
+                    "watermark",
+                    default=True,
+                    tooltip="Whether to add an \"AI generated\" watermark to the image",
+                    optional=True,
+                ),
+            ],
+            outputs=[
+                comfy_io.Image.Output(),
+            ],
+            hidden=[
+                comfy_io.Hidden.auth_token_comfy_org,
+                comfy_io.Hidden.api_key_comfy_org,
+                comfy_io.Hidden.unique_id,
+            ],
+            is_api_node=True,
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        model: str,
+        image: torch.Tensor,
+        prompt: str,
+        seed: int,
+        guidance_scale: float,
+        watermark: bool,
+    ) -> comfy_io.NodeOutput:
+        validate_string(prompt, strip_whitespace=True, min_length=1)
+        if get_number_of_images(image) != 1:
+            raise ValueError("Exactly one input image is required.")
+        validate_image_aspect_ratio_range(image, (1, 3), (3, 1))
+        auth_kwargs = {
+            "auth_token": cls.hidden.auth_token_comfy_org,
+            "comfy_api_key": cls.hidden.api_key_comfy_org,
+        }
+        source_url = (await upload_images_to_comfyapi(
+            image,
+            max_images=1,
+            mime_type="image/png",
+            auth_kwargs=auth_kwargs,
+        ))[0]
+        payload = Image2ImageTaskCreationRequest(
+            model=model,
+            prompt=prompt,
+            image=source_url,
+            seed=seed,
+            guidance_scale=guidance_scale,
+            watermark=watermark,
+        )
+        response = await SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path=BYTEPLUS_ENDPOINT,
+                method=HttpMethod.POST,
+                request_model=Image2ImageTaskCreationRequest,
+                response_model=ImageTaskCreationResponse,
+            ),
+            request=payload,
+            auth_kwargs=auth_kwargs,
+        ).execute()
+        return comfy_io.NodeOutput(await download_url_to_image_tensor(get_image_url_from_response(response)))
+
+
+class ByteDanceExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[comfy_io.ComfyNode]]:
+        return [
+            ByteDanceImageNode,
+            ByteDanceImageEditNode,
+        ]
+
+async def comfy_entrypoint() -> ByteDanceExtension:
+    return ByteDanceExtension()

nodes.py

@@ -2341,6 +2341,7 @@ async def init_builtin_api_nodes():
         "nodes_veo2.py",
         "nodes_kling.py",
         "nodes_bfl.py",
+        "nodes_bytedance.py",
         "nodes_luma.py",
         "nodes_recraft.py",
         "nodes_pixverse.py",