diff --git a/comfy_api_nodes/apis/qwen.py b/comfy_api_nodes/apis/qwen.py new file mode 100644 index 000000000..90b68dee8 --- /dev/null +++ b/comfy_api_nodes/apis/qwen.py @@ -0,0 +1,46 @@ +from pydantic import BaseModel, Field + + +class QwenImageContentItem(BaseModel): + image: str | None = Field(None) + text: str | None = Field(None) + + +class QwenImageMessage(BaseModel): + role: str = Field("user") + content: list[QwenImageContentItem] = Field(...) + + +class QwenImageInputField(BaseModel): + messages: list[QwenImageMessage] = Field(...) + + +class QwenImageParametersField(BaseModel): + size: str | None = Field(None, description="Output resolution as 'width*height'; omit for the model default.") + n: int = Field(1, ge=1, le=6) + seed: int = Field(..., ge=0, le=2147483647) + prompt_extend: bool = Field(True) + watermark: bool = Field(False) + negative_prompt: str | None = Field(None) + + +class QwenImageGenerationRequest(BaseModel): + model: str = Field(...) + input: QwenImageInputField = Field(...) + parameters: QwenImageParametersField = Field(...) + + +class QwenImageChoice(BaseModel): + finish_reason: str | None = Field(None) + message: QwenImageMessage | None = Field(None) + + +class QwenImageOutputField(BaseModel): + choices: list[QwenImageChoice] = Field(default_factory=list) + + +class QwenImageGenerationResponse(BaseModel): + output: QwenImageOutputField | None = Field(None) + request_id: str = Field(...) + code: str | None = Field(None, description="Error code for the failed request.") + message: str | None = Field(None, description="Details about the failed request.") diff --git a/comfy_api_nodes/nodes_qwen.py b/comfy_api_nodes/nodes_qwen.py new file mode 100644 index 000000000..3b6c5023c --- /dev/null +++ b/comfy_api_nodes/nodes_qwen.py @@ -0,0 +1,442 @@ +import math +import re + +import torch +from typing_extensions import override + +from comfy_api.latest import IO, ComfyExtension +from comfy_api_nodes.apis.qwen import ( + QwenImageContentItem, + QwenImageGenerationRequest, + QwenImageGenerationResponse, + QwenImageInputField, + QwenImageMessage, + QwenImageParametersField, +) +from comfy_api_nodes.util import ( + ApiEndpoint, + download_url_to_image_tensor, + sync_op, + tensor_to_base64_string, + validate_string, +) + +GENERATION_PATH = "/proxy/qwen/api/v1/services/aigc/multimodal-generation/generation" +QWEN_IMAGE_MODELS = ["qwen-image-3.0-pro", "qwen-image-3.0"] +MIN_AREA = 262144 # 512*512 +MAX_AREA = 6553600 # 2560*2560 +MAX_ASPECT = 8 # the API allows aspect ratios from 1:8 to 8:1 +MAX_INPUT_BYTES = 10 * 1024 * 1024 # the API rejects decoded input images over 10MB + +_IMAGE_REF_RE = re.compile(r"@image(?P\d*)(?!\w)", re.IGNORECASE | re.ASCII) + + +def _resolve_image_refs(prompt: str, total_images: int) -> str: + """Rewrite @Image1-style references (shared partner-node syntax, 1-based; an unnumbered + @image means the first image) into the plain 'Image N' wording the model resolves + natively. A tag counts only at a word boundary or right after a previous tag, so + adjacent tags like '@Image1@Image2' all resolve while addresses like user@image1.com + pass through untouched.""" + parts = [] + pos = 0 + prev_end = -1 + for match in _IMAGE_REF_RE.finditer(prompt): + start = match.start() + if start > 0 and start != prev_end and (prompt[start - 1].isalnum() or prompt[start - 1] == "_"): + continue + idx = int(match.group("idx") or 1) + if not 1 <= idx <= total_images: + raise ValueError( + f"The prompt references @Image{idx}, but only {total_images} reference images " + f"are connected (a batched input counts once per image)." + ) + parts.append(prompt[pos:start]) + parts.append(f"Image {idx}") + pos = match.end() + prev_end = match.end() + parts.append(prompt[pos:]) + return "".join(parts) + + +def _validate_size(width: int, height: int) -> None: + if not MIN_AREA <= width * height <= MAX_AREA: + raise ValueError( + f"Image area must be between {MIN_AREA} (512x512) and {MAX_AREA} (2560x2560) pixels; " + f"got {width}x{height} = {width * height}." + ) + if width > MAX_ASPECT * height or height > MAX_ASPECT * width: + raise ValueError(f"Aspect ratio must be between 1:8 and 8:1; got {width}x{height}.") + + +def _fit_to_size(width: int, height: int) -> tuple[int, int]: + """Scale dimensions into the supported pixel area and 1:8..8:1 aspect range, preserving + the aspect ratio where possible.""" + if width > MAX_ASPECT * height: + height = math.ceil(width / MAX_ASPECT) + elif height > MAX_ASPECT * width: + width = math.ceil(height / MAX_ASPECT) + area = width * height + if area < MIN_AREA: + scale = math.sqrt(MIN_AREA / area) + width, height = math.ceil(width * scale), math.ceil(height * scale) + elif area > MAX_AREA: + scale = math.sqrt(MAX_AREA / area) + width, height = math.floor(width * scale), math.floor(height * scale) + # rounding can push the ratio a hair past the limit; trimming only ever shrinks the area + return min(width, MAX_ASPECT * height), min(height, MAX_ASPECT * width) + + +def _image_data_uri(image: torch.Tensor) -> str: + """PNG data URI of an RGB view of the image, downscaled to <=2048x2048; falls back to + JPEG when the PNG exceeds the API's decoded-size cap (e.g. noisy, incompressible images).""" + image = image[..., :3] + b64 = tensor_to_base64_string(image, total_pixels=2048 * 2048) + if len(b64) * 3 > MAX_INPUT_BYTES * 4: + return "data:image/jpeg;base64," + tensor_to_base64_string( + image, total_pixels=2048 * 2048, mime_type="image/jpeg" + ) + return "data:image/png;base64," + b64 + + +async def _download_result_images(response: QwenImageGenerationResponse) -> torch.Tensor: + if not response.output: + raise Exception(f"An unknown error occurred: {response.code} - {response.message}") + urls = [ + item.image + for choice in response.output.choices + if choice.message + for item in choice.message.content + if item.image + ] + if not urls: + raise Exception(f"The response contains no images: {response.code} - {response.message}") + return torch.cat([await download_url_to_image_tensor(url) for url in urls]) + + +def _size_inputs() -> list[IO.Int.Input]: + return [ + IO.Int.Input( + "width", + default=1024, + min=256, + max=2560, + step=16, + tooltip="The total pixel area must be between 512x512 and 2560x2560; " + "any aspect ratio within that area works.", + ), + IO.Int.Input( + "height", + default=1024, + min=256, + max=2560, + step=16, + tooltip="The total pixel area must be between 512x512 and 2560x2560; " + "any aspect ratio within that area works.", + ), + ] + + +def _t2i_model_option(model_id: str) -> IO.DynamicCombo.Option: + return IO.DynamicCombo.Option( + model_id, + [ + IO.String.Input( + "prompt", + multiline=True, + default="", + tooltip="Prompt describing the image. Supports English and Chinese.", + ), + IO.String.Input( + "negative_prompt", + multiline=True, + default="", + tooltip="Negative prompt describing what to avoid.", + ), + *_size_inputs(), + ], + ) + + +def _edit_model_option(model_id: str) -> IO.DynamicCombo.Option: + return IO.DynamicCombo.Option( + model_id, + [ + IO.Autogrow.Input( + "images", + template=IO.Autogrow.TemplateNames( + IO.Image.Input("image"), + names=["image_1", "image_2", "image_3"], + min=1, + ), + tooltip="1-3 reference images. Refer to them in the prompt as @Image1, @Image2, " + "@Image3, numbered in input order; a batched input counts once per image.", + ), + IO.String.Input( + "prompt", + multiline=True, + default="", + tooltip="Editing instructions. Supports English and Chinese, " + "and @Image1-style references to the input images.", + ), + IO.String.Input( + "negative_prompt", + multiline=True, + default="", + tooltip="Negative prompt describing what to avoid.", + ), + ], + ) + + +class QwenImageTextToImageApi(IO.ComfyNode): + @classmethod + def define_schema(cls): + return IO.Schema( + node_id="QwenImageTextToImageApi", + display_name="Qwen Image 3 Text to Image", + category="partner/image/Qwen", + description="Generates images from a text prompt using the Qwen-Image 3.0 models.", + inputs=[ + IO.DynamicCombo.Input( + "model", + options=[_t2i_model_option(model_id) for model_id in QWEN_IMAGE_MODELS], + tooltip="Model to use.", + ), + IO.Int.Input( + "n", + default=1, + min=1, + max=6, + display_mode=IO.NumberDisplay.number, + tooltip="Number of images to generate, returned as a batch.", + ), + IO.Int.Input( + "seed", + default=42, + min=0, + max=2147483647, + step=1, + display_mode=IO.NumberDisplay.number, + control_after_generate=True, + tooltip="Seed to use for generation.", + ), + IO.Boolean.Input( + "prompt_extend", + default=True, + tooltip="Whether to enhance the prompt with AI assistance.", + advanced=True, + ), + IO.Boolean.Input( + "watermark", + default=False, + tooltip="Whether to add an AI-generated watermark to the result.", + advanced=True, + ), + ], + 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, + price_badge=IO.PriceBadge( + depends_on=IO.PriceBadgeDepends(widgets=["model", "model.width", "model.height", "n"]), + expr=""" + ( + $isPro := widgets.model = "qwen-image-3.0-pro"; + $area := $lookup(widgets, "model.width") * $lookup(widgets, "model.height"); + $rate := $isPro ? ($area > 2250000 ? 0.10725 : 0.0572) : 0.0429; + {"type":"usd","usd": $rate * widgets.n} + ) + """, + ), + ) + + @classmethod + async def execute( + cls, + model: dict, + n: int = 1, + seed: int = 42, + prompt_extend: bool = True, + watermark: bool = False, + ): + validate_string(model["prompt"], strip_whitespace=False, min_length=1) + width, height = model["width"], model["height"] + _validate_size(width, height) + response = await sync_op( + cls, + ApiEndpoint(path=GENERATION_PATH, method="POST"), + response_model=QwenImageGenerationResponse, + data=QwenImageGenerationRequest( + model=model["model"], + input=QwenImageInputField( + messages=[QwenImageMessage(content=[QwenImageContentItem(text=model["prompt"])])], + ), + parameters=QwenImageParametersField( + size=f"{width}*{height}", + n=n, + seed=seed, + prompt_extend=prompt_extend, + watermark=watermark, + negative_prompt=model["negative_prompt"] or None, + ), + ), + ) + return IO.NodeOutput(await _download_result_images(response)) + + +class QwenImageEditApi(IO.ComfyNode): + @classmethod + def define_schema(cls): + return IO.Schema( + node_id="QwenImageEditApi", + display_name="Qwen Image 3 Edit", + category="partner/image/Qwen", + description="Edits or combines up to 3 reference images guided by a text prompt " + "using the Qwen-Image 3.0 models.", + inputs=[ + IO.DynamicCombo.Input( + "model", + options=[_edit_model_option(model_id) for model_id in QWEN_IMAGE_MODELS], + tooltip="Model to use.", + ), + IO.DynamicCombo.Input( + "size", + options=[ + IO.DynamicCombo.Option("match input", []), + IO.DynamicCombo.Option("auto", []), + IO.DynamicCombo.Option("custom", _size_inputs()), + ], + tooltip="Output resolution. 'match input' reuses the first reference image's size, " + "'auto' lets the model pick a size with the same aspect ratio, " + "'custom' sets an explicit width and height.", + ), + IO.Int.Input( + "n", + default=1, + min=1, + max=6, + display_mode=IO.NumberDisplay.number, + tooltip="Number of images to generate, returned as a batch.", + ), + IO.Int.Input( + "seed", + default=42, + min=0, + max=2147483647, + step=1, + display_mode=IO.NumberDisplay.number, + control_after_generate=True, + tooltip="Seed to use for generation.", + ), + IO.Boolean.Input( + "prompt_extend", + default=True, + tooltip="Whether to enhance the prompt with AI assistance.", + advanced=True, + ), + IO.Boolean.Input( + "watermark", + default=False, + tooltip="Whether to add an AI-generated watermark to the result.", + advanced=True, + ), + ], + 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, + price_badge=IO.PriceBadge( + depends_on=IO.PriceBadgeDepends( + widgets=["model", "size", "size.width", "size.height", "n"], + input_groups=["model.images"], + ), + expr=""" + ( + $isPro := widgets.model = "qwen-image-3.0-pro"; + $mode := widgets.size; + $count := $max([$lookup(inputGroups, "model.images"), 1]); + $inputCost := 0.00429 * $count; + $area := $mode = "custom" + ? $lookup(widgets, "size.width") * $lookup(widgets, "size.height") : 0; + $customRate := $area > 2250000 ? 0.10725 : 0.0572; + $isPro and $mode != "custom" + ? {"type":"range_usd", + "min_usd": 0.0572 * widgets.n + $inputCost, + "max_usd": 0.10725 * widgets.n + $inputCost} + : {"type":"usd", + "usd": ($isPro ? $customRate : 0.0429) * widgets.n + $inputCost} + ) + """, + ), + ) + + @classmethod + async def execute( + cls, + model: dict, + size: dict, + n: int = 1, + seed: int = 42, + prompt_extend: bool = True, + watermark: bool = False, + ): + validate_string(model["prompt"], strip_whitespace=False, min_length=1) + reference_images = [image for key in model["images"] for image in model["images"][key]] + if len(reference_images) > 3: + raise ValueError( + f"A maximum of 3 reference images is supported; got {len(reference_images)} " + f"(a batched input counts once per image)." + ) + prompt = _resolve_image_refs(model["prompt"], len(reference_images)) + if size["size"] == "custom": + _validate_size(size["width"], size["height"]) + size_str = f"{size['width']}*{size['height']}" + elif size["size"] == "match input": + height, width = reference_images[0].shape[0], reference_images[0].shape[1] + width, height = _fit_to_size(width, height) + size_str = f"{width}*{height}" + else: # auto: the API picks a size preserving the input aspect ratio (1.9-4.2 MP) + size_str = None + content = [QwenImageContentItem(image=_image_data_uri(image)) for image in reference_images] + content.append(QwenImageContentItem(text=prompt)) + response = await sync_op( + cls, + ApiEndpoint(path=GENERATION_PATH, method="POST"), + response_model=QwenImageGenerationResponse, + data=QwenImageGenerationRequest( + model=model["model"], + input=QwenImageInputField(messages=[QwenImageMessage(content=content)]), + parameters=QwenImageParametersField( + size=size_str, + n=n, + seed=seed, + prompt_extend=prompt_extend, + watermark=watermark, + negative_prompt=model["negative_prompt"] or None, + ), + ), + ) + return IO.NodeOutput(await _download_result_images(response)) + + +class QwenApiExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[IO.ComfyNode]]: + return [ + QwenImageTextToImageApi, + QwenImageEditApi, + ] + + +async def comfy_entrypoint() -> QwenApiExtension: + return QwenApiExtension()