comfy_api_nodes/nodes_bfl.py

@@ -1,8 +1,10 @@
+from inspect import cleandoc
+
 import torch
 from pydantic import BaseModel
 from typing_extensions import override
 
-from comfy_api.latest import IO, ComfyExtension, Input
+from comfy_api.latest import IO, ComfyExtension
 from comfy_api_nodes.apis.bfl_api import (
     BFLFluxExpandImageRequest,
     BFLFluxFillImageRequest,
@@ -26,7 +28,7 @@ from comfy_api_nodes.util import (
 )
 
 
-def convert_mask_to_image(mask: Input.Image):
+def convert_mask_to_image(mask: torch.Tensor):
     """
     Make mask have the expected amount of dims (4) and channels (3) to be recognized as an image.
     """
@@ -36,6 +38,9 @@ def convert_mask_to_image(mask: Input.Image):
 
 
 class FluxProUltraImageNode(IO.ComfyNode):
+    """
+    Generates images using Flux Pro 1.1 Ultra via api based on prompt and resolution.
+    """
 
     @classmethod
     def define_schema(cls) -> IO.Schema:
@@ -43,7 +48,7 @@ class FluxProUltraImageNode(IO.ComfyNode):
             node_id="FluxProUltraImageNode",
             display_name="Flux 1.1 [pro] Ultra Image",
             category="api node/image/BFL",
-            description="Generates images using Flux Pro 1.1 Ultra via api based on prompt and resolution.",
+            description=cleandoc(cls.__doc__ or ""),
             inputs=[
                 IO.String.Input(
                     "prompt",
@@ -112,7 +117,7 @@ class FluxProUltraImageNode(IO.ComfyNode):
         prompt_upsampling: bool = False,
         raw: bool = False,
         seed: int = 0,
-        image_prompt: Input.Image | None = None,
+        image_prompt: torch.Tensor | None = None,
         image_prompt_strength: float = 0.1,
     ) -> IO.NodeOutput:
         if image_prompt is None:
@@ -150,6 +155,9 @@ class FluxProUltraImageNode(IO.ComfyNode):
 
 
 class FluxKontextProImageNode(IO.ComfyNode):
+    """
+    Edits images using Flux.1 Kontext [pro] via api based on prompt and aspect ratio.
+    """
 
     @classmethod
     def define_schema(cls) -> IO.Schema:
@@ -157,7 +165,7 @@ class FluxKontextProImageNode(IO.ComfyNode):
             node_id=cls.NODE_ID,
             display_name=cls.DISPLAY_NAME,
             category="api node/image/BFL",
-            description="Edits images using Flux.1 Kontext [pro] via api based on prompt and aspect ratio.",
+            description=cleandoc(cls.__doc__ or ""),
             inputs=[
                 IO.String.Input(
                     "prompt",
@@ -223,7 +231,7 @@ class FluxKontextProImageNode(IO.ComfyNode):
         aspect_ratio: str,
         guidance: float,
         steps: int,
-        input_image: Input.Image | None = None,
+        input_image: torch.Tensor | None = None,
         seed=0,
         prompt_upsampling=False,
     ) -> IO.NodeOutput:
@@ -263,14 +271,20 @@ class FluxKontextProImageNode(IO.ComfyNode):
 
 
 class FluxKontextMaxImageNode(FluxKontextProImageNode):
+    """
+    Edits images using Flux.1 Kontext [max] via api based on prompt and aspect ratio.
+    """
 
-    DESCRIPTION = "Edits images using Flux.1 Kontext [max] via api based on prompt and aspect ratio."
+    DESCRIPTION = cleandoc(__doc__ or "")
     BFL_PATH = "/proxy/bfl/flux-kontext-max/generate"
     NODE_ID = "FluxKontextMaxImageNode"
     DISPLAY_NAME = "Flux.1 Kontext [max] Image"
 
 
 class FluxProExpandNode(IO.ComfyNode):
+    """
+    Outpaints image based on prompt.
+    """
 
     @classmethod
     def define_schema(cls) -> IO.Schema:
@@ -278,7 +292,7 @@ class FluxProExpandNode(IO.ComfyNode):
             node_id="FluxProExpandNode",
             display_name="Flux.1 Expand Image",
             category="api node/image/BFL",
-            description="Outpaints image based on prompt.",
+            description=cleandoc(cls.__doc__ or ""),
             inputs=[
                 IO.Image.Input("image"),
                 IO.String.Input(
@@ -357,7 +371,7 @@ class FluxProExpandNode(IO.ComfyNode):
     @classmethod
     async def execute(
         cls,
-        image: Input.Image,
+        image: torch.Tensor,
         prompt: str,
         prompt_upsampling: bool,
         top: int,
@@ -404,6 +418,9 @@ class FluxProExpandNode(IO.ComfyNode):
 
 
 class FluxProFillNode(IO.ComfyNode):
+    """
+    Inpaints image based on mask and prompt.
+    """
 
     @classmethod
     def define_schema(cls) -> IO.Schema:
@@ -411,7 +428,7 @@ class FluxProFillNode(IO.ComfyNode):
             node_id="FluxProFillNode",
             display_name="Flux.1 Fill Image",
             category="api node/image/BFL",
-            description="Inpaints image based on mask and prompt.",
+            description=cleandoc(cls.__doc__ or ""),
             inputs=[
                 IO.Image.Input("image"),
                 IO.Mask.Input("mask"),
@@ -463,8 +480,8 @@ class FluxProFillNode(IO.ComfyNode):
     @classmethod
     async def execute(
         cls,
-        image: Input.Image,
-        mask: Input.Image,
+        image: torch.Tensor,
+        mask: torch.Tensor,
         prompt: str,
         prompt_upsampling: bool,
         steps: int,
@@ -508,15 +525,11 @@ class FluxProFillNode(IO.ComfyNode):
 
 class Flux2ProImageNode(IO.ComfyNode):
 
-    NODE_ID = "Flux2ProImageNode"
-    DISPLAY_NAME = "Flux.2 [pro] Image"
-    API_ENDPOINT = "/proxy/bfl/flux-2-pro/generate"
-
     @classmethod
     def define_schema(cls) -> IO.Schema:
         return IO.Schema(
-            node_id=cls.NODE_ID,
-            display_name=cls.DISPLAY_NAME,
+            node_id="Flux2ProImageNode",
+            display_name="Flux.2 [pro] Image",
             category="api node/image/BFL",
             description="Generates images synchronously based on prompt and resolution.",
             inputs=[
@@ -550,11 +563,12 @@ class Flux2ProImageNode(IO.ComfyNode):
                 ),
                 IO.Boolean.Input(
                     "prompt_upsampling",
-                    default=True,
+                    default=False,
                     tooltip="Whether to perform upsampling on the prompt. "
-                    "If active, automatically modifies the prompt for more creative generation.",
+                    "If active, automatically modifies the prompt for more creative generation, "
+                    "but results are nondeterministic (same seed will not produce exactly the same result).",
                 ),
-                IO.Image.Input("images", optional=True, tooltip="Up to 9 images to be used as references."),
+                IO.Image.Input("images", optional=True, tooltip="Up to 4 images to be used as references."),
             ],
             outputs=[IO.Image.Output()],
             hidden=[
@@ -573,7 +587,7 @@ class Flux2ProImageNode(IO.ComfyNode):
         height: int,
         seed: int,
         prompt_upsampling: bool,
-        images: Input.Image | None = None,
+        images: torch.Tensor | None = None,
     ) -> IO.NodeOutput:
         reference_images = {}
         if images is not None:
@@ -584,7 +598,7 @@ class Flux2ProImageNode(IO.ComfyNode):
                 reference_images[key_name] = tensor_to_base64_string(images[image_index], total_pixels=2048 * 2048)
         initial_response = await sync_op(
             cls,
-            ApiEndpoint(path=cls.API_ENDPOINT, method="POST"),
+            ApiEndpoint(path="/proxy/bfl/flux-2-pro/generate", method="POST"),
             response_model=BFLFluxProGenerateResponse,
             data=Flux2ProGenerateRequest(
                 prompt=prompt,
@@ -618,13 +632,6 @@ class Flux2ProImageNode(IO.ComfyNode):
         return IO.NodeOutput(await download_url_to_image_tensor(response.result["sample"]))
 
 
-class Flux2MaxImageNode(Flux2ProImageNode):
-
-    NODE_ID = "Flux2MaxImageNode"
-    DISPLAY_NAME = "Flux.2 [max] Image"
-    API_ENDPOINT = "/proxy/bfl/flux-2-max/generate"
-
-
 class BFLExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[IO.ComfyNode]]:
@@ -635,7 +642,6 @@ class BFLExtension(ComfyExtension):
             FluxProExpandNode,
             FluxProFillNode,
             Flux2ProImageNode,
-            Flux2MaxImageNode,
         ]
 
 

manager_requirements.txt

@@ -1 +1 @@
-comfyui_manager==4.0.3b7
+comfyui_manager==4.0.3b5