Move and fix BFL nodes to node_bfl.py (#49)

comfy_api_nodes/apis/BFLPolling.py

@@ -1,29 +0,0 @@
-from __future__ import annotations
-
-from enum import Enum
-from typing import Any, Dict, Optional
-
-from pydantic import BaseModel, Field, confloat
-
-
-class BFLStatus(str, Enum):
-    task_not_found = "Task not found"
-    pending = "Pending"
-    request_moderated = "Request Moderated"
-    content_moderated = "Content Moderated"
-    ready = "Ready"
-    error = "Error"
-
-
-class BFLFluxProStatusResponse(BaseModel):
-    id: str = Field(..., description="The unique identifier for the generation task.")
-    status: BFLStatus = Field(..., description="The status of the task.")
-    result: Optional[Dict[str, Any]] = Field(
-        None, description="The result of the task (null if not completed)."
-    )
-    progress: confloat(ge=0.0, le=1.0) = Field(
-        ..., description="The progress of the task (0.0 to 1.0)."
-    )
-    details: Optional[Dict[str, Any]] = Field(
-        None, description="Additional details about the task (null if not available)."
-    )

comfy_api_nodes/apis/bfl_api.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, Field, confloat, conint
+
+
+class BFLOutputFormat(str, Enum):
+    png = 'png'
+    jpeg = 'jpeg'
+
+
+class BFLFluxProGenerateRequest(BaseModel):
+    prompt: str = Field(..., description='The text prompt for image generation.')
+    prompt_upsampling: Optional[bool] = Field(
+        None, description='Whether to perform upsampling on the prompt. If active, automatically modifies the prompt for more creative generation.'
+    )
+    seed: Optional[int] = Field(None, description='The seed value for reproducibility.')
+    aspect_ratio: Optional[str] = Field(None, description='Aspect ratio of the image between 21:9 and 9:21.')
+    safety_tolerance: Optional[conint(ge=0, le=6)] = Field(
+        6, description='Tolerance level for input and output moderation. Between 0 and 6, 0 being most strict, 6 being least strict. Defaults to 2.'
+    )
+    output_format: Optional[BFLOutputFormat] = Field(
+        BFLOutputFormat.png, description="Output format for the generated image. Can be 'jpeg' or 'png'.", examples=['png']
+    )
+    raw: Optional[bool] = Field(None, description='Generate less processed, more natural-looking images.')
+    image_prompt: Optional[str] = Field(None, description='Optional image to remix in base64 format')
+    image_prompt_strength: Optional[confloat(ge=0.0, le=1.0)] = Field(
+        None, description='Blend between the prompt and the image prompt.'
+    )
+
+
+class BFLFluxProGenerateResponse(BaseModel):
+    id: str = Field(..., description='The unique identifier for the generation task.')
+    polling_url: str = Field(..., description='URL to poll for the generation result.')
+
+
+class BFLStatus(str, Enum):
+    task_not_found = "Task not found"
+    pending = "Pending"
+    request_moderated = "Request Moderated"
+    content_moderated = "Content Moderated"
+    ready = "Ready"
+    error = "Error"
+
+
+class BFLFluxProStatusResponse(BaseModel):
+    id: str = Field(..., description="The unique identifier for the generation task.")
+    status: BFLStatus = Field(..., description="The status of the task.")
+    result: Optional[Dict[str, Any]] = Field(
+        None, description="The result of the task (null if not completed)."
+    )
+    progress: confloat(ge=0.0, le=1.0) = Field(
+        ..., description="The progress of the task (0.0 to 1.0)."
+    )
+    details: Optional[Dict[str, Any]] = Field(
+        None, description="Additional details about the task (null if not available)."
+    )

comfy_api_nodes/nodes_api.py

@@ -17,12 +17,9 @@ from comfy_api_nodes.apis import (
     MinimaxTaskResultResponse,
     IdeogramGenerateRequest,
     IdeogramGenerateResponse,
-    BFLFluxProGenerateRequest,
-    BFLFluxProGenerateResponse,
     ImageRequest,
     Model
 )
-from comfy_api_nodes.apis.BFLPolling import BFLStatus
 from comfy_api_nodes.apis.client import (
     ApiClient,
     ApiEndpoint,
@@ -41,7 +38,6 @@ import torch
 import math
 import base64
 import logging
-import time
 import uuid
 import folder_paths
 from io import BytesIO
@@ -1000,212 +996,6 @@ class IdeogramTextToImage(ComfyNodeABC):
     # def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen):
     #    return ""
 
-
-class FluxProUltraImageNode(ComfyNodeABC):
-    """
-    Generates images synchronously based on prompt and resolution.
-    """
-
-    MINIMUM_RATIO = 1 / 4
-    MAXIMUM_RATIO = 4 / 1
-    MINIMUM_RATIO_STR = "1:4"
-    MAXIMUM_RATIO_STR = "4:1"
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "prompt": (
-                    IO.STRING,
-                    {
-                        "multiline": True,
-                        "default": "",
-                        "tooltip": "Prompt for the image generation",
-                    },
-                ),
-                "prompt_upsampling": (
-                    IO.BOOLEAN,
-                    {
-                        "default": False,
-                        "tooltip": "Whether to perform upsampling on the prompt. If active, automatically modifies the prompt for more creative generation, but results are nondeterministic (same seed will not produce exactly the same result).",
-                    },
-                ),
-                "seed": (
-                    IO.INT,
-                    {
-                        "default": 0,
-                        "min": 0,
-                        "max": 0xFFFFFFFFFFFFFFFF,
-                        "control_after_generate": True,
-                        "tooltip": "The random seed used for creating the noise.",
-                    },
-                ),
-                "aspect_ratio": (
-                    IO.STRING,
-                    {
-                        "default": "16:9",
-                        "tooltip": "Aspect ratio of image; must be between 1:4 and 4:1.",
-                    },
-                ),
-                "raw": (
-                    IO.BOOLEAN,
-                    {
-                        "default": False,
-                        "tooltip": "When True, generate less processed, more natural-looking images.",
-                    },
-                ),
-            },
-            "optional": {
-                "image_prompt": (IO.IMAGE,),
-                "image_prompt_strength": (
-                    IO.FLOAT,
-                    {
-                        "default": 0.1,
-                        "min": 0.0,
-                        "max": 1.0,
-                        "step": 0.01,
-                        "tooltip": "Blend between the prompt and the image prompt.",
-                    },
-                ),
-            },
-            "hidden": {
-                "auth_token": "AUTH_TOKEN_COMFY_ORG",
-            },
-        }
-
-    @classmethod
-    def VALIDATE_INPUTS(cls, aspect_ratio: str):
-        try:
-            validate_aspect_ratio(
-                aspect_ratio,
-                minimum_ratio=cls.MINIMUM_RATIO,
-                maximum_ratio=cls.MAXIMUM_RATIO,
-                minimum_ratio_str=cls.MINIMUM_RATIO_STR,
-                maximum_ratio_str=cls.MAXIMUM_RATIO_STR,
-            )
-        except Exception as e:
-            return str(e)
-        return True
-
-    RETURN_TYPES = (IO.IMAGE,)
-    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
-    FUNCTION = "api_call"
-    API_NODE = True
-    CATEGORY = "api node/image/bfl"
-
-    def api_call(
-        self,
-        prompt: str,
-        aspect_ratio: str,
-        prompt_upsampling=False,
-        raw=False,
-        seed=0,
-        image_prompt=None,
-        image_prompt_strength=0.1,
-        auth_token=None,
-        **kwargs,
-    ):
-        operation = SynchronousOperation(
-            endpoint=ApiEndpoint(
-                path="/proxy/bfl/flux-pro-1.1-ultra/generate",
-                method=HttpMethod.POST,
-                request_model=BFLFluxProGenerateRequest,
-                response_model=BFLFluxProGenerateResponse,
-            ),
-            request=BFLFluxProGenerateRequest(
-                prompt=prompt,
-                prompt_upsampling=prompt_upsampling,
-                seed=seed,
-                aspect_ratio=validate_aspect_ratio(
-                    aspect_ratio,
-                    minimum_ratio=self.MINIMUM_RATIO,
-                    maximum_ratio=self.MAXIMUM_RATIO,
-                    minimum_ratio_str=self.MINIMUM_RATIO_STR,
-                    maximum_ratio_str=self.MAXIMUM_RATIO_STR,
-                ),
-                raw=raw,
-                image_prompt=(
-                    image_prompt
-                    if image_prompt is None
-                    else self._convert_image_to_base64(image_prompt)
-                ),
-                image_prompt_strength=(
-                    None if image_prompt is None else round(image_prompt_strength, 2)
-                ),
-            ),
-            auth_token=auth_token,
-        )
-        output_image = self._handle_bfl_synchronous_operation(operation)
-        return (output_image,)
-
-    def _handle_bfl_synchronous_operation(
-        self, operation: SynchronousOperation, timeout_bfl_calls=360
-    ):
-        response_api: BFLFluxProGenerateResponse = operation.execute()
-        return self._poll_until_generated(
-            response_api.polling_url, timeout=timeout_bfl_calls
-        )
-
-    def _poll_until_generated(self, polling_url: str, timeout=360):
-        # used bfl-comfy-nodes to verify code implementation:
-        # https://github.com/black-forest-labs/bfl-comfy-nodes/tree/main
-        start_time = time.time()
-        retries_404 = 0
-        max_retries_404 = 5
-        retry_404_seconds = 2
-        retry_202_seconds = 2
-        retry_pending_seconds = 1
-        request = requests.Request(method=HttpMethod.GET, url=polling_url)
-        # NOTE: should True loop be replaced with checking if workflow has been interrupted?
-        while True:
-            response = requests.Session().send(request.prepare())
-            if response.status_code == 200:
-                result = response.json()
-                if result["status"] == BFLStatus.ready:
-                    img_url = result["result"]["sample"]
-                    img_response = requests.get(img_url)
-                    return process_image_response(img_response)
-                elif result["status"] in [
-                    BFLStatus.request_moderated,
-                    BFLStatus.content_moderated,
-                ]:
-                    status = result["status"]
-                    raise Exception(
-                        f"BFL API did not return an image due to: {status}."
-                    )
-                elif result["status"] == BFLStatus.error:
-                    raise Exception(f"BFL API encountered an error: {result}.")
-                elif result["status"] == BFLStatus.pending:
-                    time.sleep(retry_pending_seconds)
-                    continue
-            elif response.status_code == 404:
-                if retries_404 < max_retries_404:
-                    retries_404 += 1
-                    time.sleep(retry_404_seconds)
-                    continue
-                raise Exception(
-                    f"BFL API could not find task after {max_retries_404} tries."
-                )
-            elif response.status_code == 202:
-                time.sleep(retry_202_seconds)
-            elif time.time() - start_time > timeout:
-                raise Exception(
-                    f"BFL API experienced a timeout; could not return request under {timeout} seconds."
-                )
-            else:
-                raise Exception(f"BFL API encountered an error: {response.json()}")
-
-    def _convert_image_to_base64(self, image: torch.Tensor):
-        scaled_image = downscale_input(image, total_pixels=2048 * 2048)
-        # remove batch dimension if present
-        if len(scaled_image.shape) > 3:
-            scaled_image = scaled_image[0]
-        image_np = (scaled_image.numpy() * 255).astype(np.uint8)
-        img = Image.fromarray(image_np)
-        img_byte_arr = io.BytesIO()
-        img.save(img_byte_arr, format="PNG")
-        return base64.b64encode(img_byte_arr.getvalue()).decode()
-
 class MinimaxTextToVideoNode:
     """
     Generates videos synchronously based on a prompt, and optional parameters using Minimax's API.
@@ -1352,7 +1142,6 @@ NODE_CLASS_MAPPINGS = {
     "OpenAIDalle3": OpenAIDalle3,
     "OpenAIGPTImage1": OpenAIGPTImage1,
     "IdeogramTextToImage": IdeogramTextToImage,
-    "FluxProUltraImageNode": FluxProUltraImageNode,
     "MinimaxTextToVideoNode": MinimaxTextToVideoNode,
 }
 
@@ -1362,6 +1151,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "OpenAIDalle3": "OpenAI DALL·E 3",
     "OpenAIGPTImage1": "OpenAI GPT Image 1",
     "IdeogramTextToImage": "Ideogram Text to Image",
-    "FluxProUltraImageNode": "Flux 1.1 [pro] Ultra Image",
     "MinimaxTextToVideoNode": "Minimax Text to Video",
 }

comfy_api_nodes/nodes_bfl.py

@@ -0,0 +1,243 @@
+import io
+from inspect import cleandoc
+from comfy.comfy_types.node_typing import IO, ComfyNodeABC
+from comfy_api_nodes.apis.bfl_api import (
+    BFLStatus,
+    BFLFluxProGenerateRequest,
+    BFLFluxProGenerateResponse,
+)
+from comfy_api_nodes.apis.client import (
+    ApiEndpoint,
+    HttpMethod,
+    SynchronousOperation,
+)
+from comfy_api_nodes.nodes_api import (
+    downscale_input,
+    validate_aspect_ratio,
+    process_image_response,
+)
+
+import numpy as np
+from PIL import Image
+import requests
+import torch
+import base64
+import time
+
+
+class FluxProUltraImageNode(ComfyNodeABC):
+    """
+    Generates images synchronously based on prompt and resolution.
+    """
+
+    MINIMUM_RATIO = 1 / 4
+    MAXIMUM_RATIO = 4 / 1
+    MINIMUM_RATIO_STR = "1:4"
+    MAXIMUM_RATIO_STR = "4:1"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt": (
+                    IO.STRING,
+                    {
+                        "multiline": True,
+                        "default": "",
+                        "tooltip": "Prompt for the image generation",
+                    },
+                ),
+                "prompt_upsampling": (
+                    IO.BOOLEAN,
+                    {
+                        "default": False,
+                        "tooltip": "Whether to perform upsampling on the prompt. If active, automatically modifies the prompt for more creative generation, but results are nondeterministic (same seed will not produce exactly the same result).",
+                    },
+                ),
+                "seed": (
+                    IO.INT,
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 0xFFFFFFFFFFFFFFFF,
+                        "control_after_generate": True,
+                        "tooltip": "The random seed used for creating the noise.",
+                    },
+                ),
+                "aspect_ratio": (
+                    IO.STRING,
+                    {
+                        "default": "16:9",
+                        "tooltip": "Aspect ratio of image; must be between 1:4 and 4:1.",
+                    },
+                ),
+                "raw": (
+                    IO.BOOLEAN,
+                    {
+                        "default": False,
+                        "tooltip": "When True, generate less processed, more natural-looking images.",
+                    },
+                ),
+            },
+            "optional": {
+                "image_prompt": (IO.IMAGE,),
+                "image_prompt_strength": (
+                    IO.FLOAT,
+                    {
+                        "default": 0.1,
+                        "min": 0.0,
+                        "max": 1.0,
+                        "step": 0.01,
+                        "tooltip": "Blend between the prompt and the image prompt.",
+                    },
+                ),
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+            },
+        }
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, aspect_ratio: str):
+        try:
+            validate_aspect_ratio(
+                aspect_ratio,
+                minimum_ratio=cls.MINIMUM_RATIO,
+                maximum_ratio=cls.MAXIMUM_RATIO,
+                minimum_ratio_str=cls.MINIMUM_RATIO_STR,
+                maximum_ratio_str=cls.MAXIMUM_RATIO_STR,
+            )
+        except Exception as e:
+            return str(e)
+        return True
+
+    RETURN_TYPES = (IO.IMAGE,)
+    DESCRIPTION = cleandoc(__doc__ or "")  # Handle potential None value
+    FUNCTION = "api_call"
+    API_NODE = True
+    CATEGORY = "api node/image/bfl"
+
+    def api_call(
+        self,
+        prompt: str,
+        aspect_ratio: str,
+        prompt_upsampling=False,
+        raw=False,
+        seed=0,
+        image_prompt=None,
+        image_prompt_strength=0.1,
+        auth_token=None,
+        **kwargs,
+    ):
+        operation = SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path="/proxy/bfl/flux-pro-1.1-ultra/generate",
+                method=HttpMethod.POST,
+                request_model=BFLFluxProGenerateRequest,
+                response_model=BFLFluxProGenerateResponse,
+            ),
+            request=BFLFluxProGenerateRequest(
+                prompt=prompt,
+                prompt_upsampling=prompt_upsampling,
+                seed=seed,
+                aspect_ratio=validate_aspect_ratio(
+                    aspect_ratio,
+                    minimum_ratio=self.MINIMUM_RATIO,
+                    maximum_ratio=self.MAXIMUM_RATIO,
+                    minimum_ratio_str=self.MINIMUM_RATIO_STR,
+                    maximum_ratio_str=self.MAXIMUM_RATIO_STR,
+                ),
+                raw=raw,
+                image_prompt=(
+                    image_prompt
+                    if image_prompt is None
+                    else self._convert_image_to_base64(image_prompt)
+                ),
+                image_prompt_strength=(
+                    None if image_prompt is None else round(image_prompt_strength, 2)
+                ),
+            ),
+            auth_token=auth_token,
+        )
+        output_image = self._handle_bfl_synchronous_operation(operation)
+        return (output_image,)
+
+    def _handle_bfl_synchronous_operation(
+        self, operation: SynchronousOperation, timeout_bfl_calls=360
+    ):
+        response_api: BFLFluxProGenerateResponse = operation.execute()
+        return self._poll_until_generated(
+            response_api.polling_url, timeout=timeout_bfl_calls
+        )
+
+    def _poll_until_generated(self, polling_url: str, timeout=360):
+        # used bfl-comfy-nodes to verify code implementation:
+        # https://github.com/black-forest-labs/bfl-comfy-nodes/tree/main
+        start_time = time.time()
+        retries_404 = 0
+        max_retries_404 = 5
+        retry_404_seconds = 2
+        retry_202_seconds = 2
+        retry_pending_seconds = 1
+        request = requests.Request(method=HttpMethod.GET, url=polling_url)
+        # NOTE: should True loop be replaced with checking if workflow has been interrupted?
+        while True:
+            response = requests.Session().send(request.prepare())
+            if response.status_code == 200:
+                result = response.json()
+                if result["status"] == BFLStatus.ready:
+                    img_url = result["result"]["sample"]
+                    img_response = requests.get(img_url)
+                    return process_image_response(img_response)
+                elif result["status"] in [
+                    BFLStatus.request_moderated,
+                    BFLStatus.content_moderated,
+                ]:
+                    status = result["status"]
+                    raise Exception(
+                        f"BFL API did not return an image due to: {status}."
+                    )
+                elif result["status"] == BFLStatus.error:
+                    raise Exception(f"BFL API encountered an error: {result}.")
+                elif result["status"] == BFLStatus.pending:
+                    time.sleep(retry_pending_seconds)
+                    continue
+            elif response.status_code == 404:
+                if retries_404 < max_retries_404:
+                    retries_404 += 1
+                    time.sleep(retry_404_seconds)
+                    continue
+                raise Exception(
+                    f"BFL API could not find task after {max_retries_404} tries."
+                )
+            elif response.status_code == 202:
+                time.sleep(retry_202_seconds)
+            elif time.time() - start_time > timeout:
+                raise Exception(
+                    f"BFL API experienced a timeout; could not return request under {timeout} seconds."
+                )
+            else:
+                raise Exception(f"BFL API encountered an error: {response.json()}")
+
+    def _convert_image_to_base64(self, image: torch.Tensor):
+        scaled_image = downscale_input(image, total_pixels=2048 * 2048)
+        # remove batch dimension if present
+        if len(scaled_image.shape) > 3:
+            scaled_image = scaled_image[0]
+        image_np = (scaled_image.numpy() * 255).astype(np.uint8)
+        img = Image.fromarray(image_np)
+        img_byte_arr = io.BytesIO()
+        img.save(img_byte_arr, format="PNG")
+        return base64.b64encode(img_byte_arr.getvalue()).decode()
+
+
+# A dictionary that contains all nodes you want to export with their names
+# NOTE: names should be globally unique
+NODE_CLASS_MAPPINGS = {
+    "FluxProUltraImageNode": FluxProUltraImageNode,
+}
+
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "FluxProUltraImageNode": "Flux 1.1 [pro] Ultra Image",
+}

nodes.py

@@ -2266,6 +2266,7 @@ def init_builtin_extra_nodes():
         "nodes_veo2.py",
         "nodes_kling.py",
         "nodes_runway.py",
+        "nodes_bfl.py",
         "nodes_luma.py",
         "nodes_recraft.py",
         "nodes_pixverse.py",