[Partner Nodes] new Flux2ImageNode and GrokImageEditNodeV2 nodes with DynamicCombo and Autogrow (#13814)
Some checks are pending
Python Linting / Run Ruff (push) Waiting to run
Python Linting / Run Pylint (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Waiting to run
Execution Tests / test (macos-latest) (push) Waiting to run
Execution Tests / test (ubuntu-latest) (push) Waiting to run
Execution Tests / test (windows-latest) (push) Waiting to run
Test server launches without errors / test (push) Waiting to run
Unit Tests / test (macos-latest) (push) Waiting to run
Unit Tests / test (ubuntu-latest) (push) Waiting to run
Unit Tests / test (windows-2022) (push) Waiting to run

This commit is contained in:
Alexander Piskun 2026-05-11 11:37:15 +03:00 committed by GitHub
parent 52976f3ea3
commit b565dc7a6c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 365 additions and 0 deletions

View File

@ -596,6 +596,7 @@ class Flux2ProImageNode(IO.ComfyNode):
depends_on=IO.PriceBadgeDepends(widgets=["width", "height"], inputs=["images"]),
expr=cls.PRICE_BADGE_EXPR,
),
is_deprecated=True,
)
@classmethod
@ -674,6 +675,175 @@ class Flux2MaxImageNode(Flux2ProImageNode):
"""
_FLUX2_MODEL_ENDPOINTS = {
"Flux.2 [pro]": "/proxy/bfl/flux-2-pro/generate",
"Flux.2 [max]": "/proxy/bfl/flux-2-max/generate",
}
def _flux2_model_inputs():
return [
IO.Int.Input(
"width",
default=1024,
min=256,
max=2048,
step=32,
),
IO.Int.Input(
"height",
default=768,
min=256,
max=2048,
step=32,
),
IO.Autogrow.Input(
"images",
template=IO.Autogrow.TemplateNames(
IO.Image.Input("image"),
names=[f"image_{i}" for i in range(1, 9)],
min=0,
),
tooltip="Optional reference image(s) for image-to-image generation. Up to 8 images.",
),
]
class Flux2ImageNode(IO.ComfyNode):
@classmethod
def define_schema(cls) -> IO.Schema:
return IO.Schema(
node_id="Flux2ImageNode",
display_name="Flux.2 Image",
category="api node/image/BFL",
description="Generate images via Flux.2 [pro] or Flux.2 [max] from a prompt and optional reference images.",
inputs=[
IO.String.Input(
"prompt",
multiline=True,
default="",
tooltip="Prompt for the image generation or edit",
),
IO.DynamicCombo.Input(
"model",
options=[
IO.DynamicCombo.Option("Flux.2 [pro]", _flux2_model_inputs()),
IO.DynamicCombo.Option("Flux.2 [max]", _flux2_model_inputs()),
],
),
IO.Int.Input(
"seed",
default=0,
min=0,
max=0xFFFFFFFFFFFFFFFF,
control_after_generate=True,
tooltip="The random seed used for creating the noise.",
),
],
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"],
input_groups=["model.images"],
),
expr="""
(
$isMax := widgets.model = "flux.2 [max]";
$MP := 1024 * 1024;
$w := $lookup(widgets, "model.width");
$h := $lookup(widgets, "model.height");
$outMP := $max([1, $floor((($w * $h) + $MP - 1) / $MP)]);
$outputCost := $isMax
? (0.07 + 0.03 * ($outMP - 1))
: (0.03 + 0.015 * ($outMP - 1));
$refMin := $isMax ? 0.03 : 0.015;
$refMax := $isMax ? 0.24 : 0.12;
$hasRefs := $lookup(inputGroups, "model.images") > 0;
$hasRefs
? {
"type": "range_usd",
"min_usd": $outputCost + $refMin,
"max_usd": $outputCost + $refMax,
"format": { "approximate": true }
}
: {"type": "usd", "usd": $outputCost}
)
""",
),
)
@classmethod
async def execute(
cls,
prompt: str,
model: dict,
seed: int,
) -> IO.NodeOutput:
model_choice = model["model"]
endpoint = _FLUX2_MODEL_ENDPOINTS[model_choice]
width = model["width"]
height = model["height"]
images_dict = model.get("images") or {}
image_tensors: list[Input.Image] = [t for t in images_dict.values() if t is not None]
n_images = sum(get_number_of_images(t) for t in image_tensors)
if n_images > 8:
raise ValueError("The current maximum number of supported images is 8.")
flat_tensors: list[torch.Tensor] = []
for tensor in image_tensors:
if len(tensor.shape) == 4:
flat_tensors.extend(tensor[i] for i in range(tensor.shape[0]))
else:
flat_tensors.append(tensor)
reference_images: dict[str, str] = {}
for idx, tensor in enumerate(flat_tensors):
key_name = f"input_image_{idx + 1}" if idx else "input_image"
reference_images[key_name] = tensor_to_base64_string(tensor, total_pixels=2048 * 2048)
initial_response = await sync_op(
cls,
ApiEndpoint(path=endpoint, method="POST"),
response_model=BFLFluxProGenerateResponse,
data=Flux2ProGenerateRequest(
prompt=prompt,
width=width,
height=height,
seed=seed,
**reference_images,
),
)
def price_extractor(_r: BaseModel) -> float | None:
return None if initial_response.cost is None else initial_response.cost / 100
response = await poll_op(
cls,
ApiEndpoint(initial_response.polling_url),
response_model=BFLFluxStatusResponse,
status_extractor=lambda r: r.status,
progress_extractor=lambda r: r.progress,
price_extractor=price_extractor,
completed_statuses=[BFLStatus.ready],
failed_statuses=[
BFLStatus.request_moderated,
BFLStatus.content_moderated,
BFLStatus.error,
BFLStatus.task_not_found,
],
queued_statuses=[],
)
return IO.NodeOutput(await download_url_to_image_tensor(response.result["sample"]))
class BFLExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
@ -685,6 +855,7 @@ class BFLExtension(ComfyExtension):
FluxProFillNode,
Flux2ProImageNode,
Flux2MaxImageNode,
Flux2ImageNode,
]

View File

@ -162,6 +162,61 @@ class GrokImageNode(IO.ComfyNode):
)
_GROK_IMAGE_EDIT_ASPECT_RATIO_OPTIONS = [
"auto",
"1:1",
"2:3",
"3:2",
"3:4",
"4:3",
"9:16",
"16:9",
"9:19.5",
"19.5:9",
"9:20",
"20:9",
"1:2",
"2:1",
]
def _grok_image_edit_model_inputs(*, max_ref_images: int, with_aspect_ratio: bool):
inputs = [
IO.Autogrow.Input(
"images",
template=IO.Autogrow.TemplateNames(
IO.Image.Input("image"),
names=[f"image_{i}" for i in range(1, max_ref_images + 1)],
min=1,
),
tooltip=(
"Reference image to edit."
if max_ref_images == 1
else f"Reference image(s) to edit. Up to {max_ref_images} images."
),
),
IO.Combo.Input("resolution", options=["1K", "2K"]),
IO.Int.Input(
"number_of_images",
default=1,
min=1,
max=10,
step=1,
tooltip="Number of edited images to generate",
display_mode=IO.NumberDisplay.number,
),
]
if with_aspect_ratio:
inputs.append(
IO.Combo.Input(
"aspect_ratio",
options=_GROK_IMAGE_EDIT_ASPECT_RATIO_OPTIONS,
tooltip="Only allowed when multiple images are connected.",
)
)
return inputs
class GrokImageEditNode(IO.ComfyNode):
@classmethod
@ -256,6 +311,7 @@ class GrokImageEditNode(IO.ComfyNode):
)
""",
),
is_deprecated=True,
)
@classmethod
@ -303,6 +359,143 @@ class GrokImageEditNode(IO.ComfyNode):
)
class GrokImageEditNodeV2(IO.ComfyNode):
@classmethod
def define_schema(cls):
return IO.Schema(
node_id="GrokImageEditNodeV2",
display_name="Grok Image Edit",
category="api node/image/Grok",
description="Modify an existing image based on a text prompt",
inputs=[
IO.String.Input(
"prompt",
multiline=True,
default="",
tooltip="The text prompt used to generate the image",
),
IO.DynamicCombo.Input(
"model",
options=[
IO.DynamicCombo.Option(
"grok-imagine-image-quality",
_grok_image_edit_model_inputs(max_ref_images=3, with_aspect_ratio=True),
),
IO.DynamicCombo.Option(
"grok-imagine-image-pro",
_grok_image_edit_model_inputs(max_ref_images=1, with_aspect_ratio=False),
),
IO.DynamicCombo.Option(
"grok-imagine-image",
_grok_image_edit_model_inputs(max_ref_images=3, with_aspect_ratio=True),
),
],
),
IO.Int.Input(
"seed",
default=0,
min=0,
max=2147483647,
step=1,
display_mode=IO.NumberDisplay.number,
control_after_generate=True,
tooltip="Seed to determine if node should re-run; "
"actual results are nondeterministic regardless of seed.",
),
],
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.resolution", "model.number_of_images"],
),
expr="""
(
$isQualityModel := widgets.model = "grok-imagine-image-quality";
$isPro := $contains(widgets.model, "pro");
$res := $lookup(widgets, "model.resolution");
$n := $lookup(widgets, "model.number_of_images");
$rate := $isQualityModel
? ($res = "1k" ? 0.05 : 0.07)
: ($isPro ? 0.07 : 0.02);
$base := $isQualityModel ? 0.01 : 0.002;
$output := $rate * $n;
$isPro
? {"type":"usd","usd": $base + $output}
: {"type":"range_usd","min_usd": $base + $output, "max_usd": 3 * $base + $output}
)
""",
),
)
@classmethod
async def execute(
cls,
prompt: str,
model: dict,
seed: int,
) -> IO.NodeOutput:
validate_string(prompt, strip_whitespace=True, min_length=1)
model_id = model["model"]
resolution = model["resolution"]
number_of_images = model["number_of_images"]
images_dict = model.get("images") or {}
aspect_ratio = model.get("aspect_ratio", "auto")
image_tensors: list[Input.Image] = [t for t in images_dict.values() if t is not None]
n_images = sum(get_number_of_images(t) for t in image_tensors)
if n_images < 1:
raise ValueError("At least one image is required for editing.")
if model_id == "grok-imagine-image-pro" and n_images > 1:
raise ValueError("The pro model supports only 1 input image.")
if model_id != "grok-imagine-image-pro" and n_images > 3:
raise ValueError("A maximum of 3 input images is supported.")
if aspect_ratio != "auto" and n_images == 1:
raise ValueError(
"Custom aspect ratio is only allowed when multiple images are connected to the image input."
)
flat_tensors: list[torch.Tensor] = []
for tensor in image_tensors:
if len(tensor.shape) == 4:
flat_tensors.extend(tensor[i] for i in range(tensor.shape[0]))
else:
flat_tensors.append(tensor)
response = await sync_op(
cls,
ApiEndpoint(path="/proxy/xai/v1/images/edits", method="POST"),
data=ImageEditRequest(
model=model_id,
images=[
InputUrlObject(url=f"data:image/png;base64,{tensor_to_base64_string(i)}") for i in flat_tensors
],
prompt=prompt,
resolution=resolution.lower(),
n=number_of_images,
seed=seed,
aspect_ratio=None if aspect_ratio == "auto" else aspect_ratio,
),
response_model=ImageGenerationResponse,
price_extractor=_extract_grok_price,
)
if len(response.data) == 1:
return IO.NodeOutput(await download_url_to_image_tensor(response.data[0].url))
return IO.NodeOutput(
torch.cat(
[await download_url_to_image_tensor(i) for i in [str(d.url) for d in response.data if d.url]],
)
)
class GrokVideoNode(IO.ComfyNode):
@classmethod
@ -737,6 +930,7 @@ class GrokExtension(ComfyExtension):
return [
GrokImageNode,
GrokImageEditNode,
GrokImageEditNodeV2,
GrokVideoNode,
GrokVideoReferenceNode,
GrokVideoEditNode,