[Partner Nodes] new Flux2ImageNode and GrokImageEditNodeV2 nodes with DynamicCombo and Autogrow (#13814)

comfy_api_nodes/nodes_bfl.py

@@ -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,
         ]
 
 

comfy_api_nodes/nodes_grok.py

@@ -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,