443 lines
16 KiB
Python
443 lines
16 KiB
Python
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<idx>\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()
|