Added initial BFL Flux 1.1 [pro] Ultra node (#11)

comfy_api_nodes/nodes_api.py

@@ -14,6 +14,9 @@ from comfy_api_nodes.apis import (
     MinimaxTaskResultResponse,
     IdeogramGenerateRequest,
     IdeogramGenerateResponse,
+    BFLFluxProGenerateRequest,
+    BFLFluxProGenerateResponse,
+    BFLStatus,
     ImageRequest,
     Model
 )
@@ -29,12 +32,13 @@ import logging
 import json
 import av
 import os
+import time
 import folder_paths
 
-def downscale_input(image):
+def downscale_input(image, total_pixels=1536*1024):
     samples = image.movedim(-1,1)
     #downscaling input images to roughly the same size as the outputs
-    total = int(1536 * 1024)
+    total = int(total_pixels)
     scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
     if scale_by >= 1:
         return image
@@ -83,6 +87,25 @@ def validate_and_cast_response(response):
 
     return torch.stack(image_tensors, dim=0)
 
+def validate_aspect_ratio(aspect_ratio: str, minimum_ratio: float, maximum_ratio: float, minimum_ratio_str: str, maximum_ratio_str: str):
+    # get ratio values
+    numbers = aspect_ratio.split(':')
+    if len(numbers) != 2:
+        raise Exception(f"Aspect ratio must be in the format X:Y, such as 16:9, but was {aspect_ratio}.")
+    try:
+        numerator = int(numbers[0])
+        denominator = int(numbers[1])
+    except ValueError:
+        raise Exception(f"Aspect ratio must contain numbers separated by ':', such as 16:9, but was {aspect_ratio}.")
+    calculated_ratio = numerator/denominator
+    # if not close to minimum and maximum, check bounds
+    if not math.isclose(calculated_ratio, minimum_ratio) or not math.isclose(calculated_ratio, maximum_ratio):
+        if calculated_ratio < minimum_ratio:
+            raise Exception(f"Aspect ratio cannot reduce to any less than {minimum_ratio_str} ({minimum_ratio}), but was {aspect_ratio} ({calculated_ratio}).")
+        elif calculated_ratio > maximum_ratio:
+            raise Exception(f"Aspect ratio cannot reduce to any greater than {maximum_ratio_str} ({maximum_ratio}), but was {aspect_ratio} ({calculated_ratio}).")
+    return aspect_ratio
+
 class OpenAIDalle2(ComfyNodeABC):
     """
     Generates images synchronously via OpenAI's DALL·E 2 endpoint.
@@ -595,6 +618,156 @@ 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"
+
+    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 self._process_bfl_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 _process_bfl_image_response(self, response: requests.Response):
+        image = Image.open(io.BytesIO(response.content)).convert("RGBA")
+        image_array = np.array(image).astype(np.float32) / 255.0
+        return torch.from_numpy(image_array).unsqueeze(0)
+
+    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:
     """
@@ -772,6 +945,7 @@ NODE_CLASS_MAPPINGS = {
     "OpenAIDalle3": OpenAIDalle3,
     "OpenAIGPTImage1": OpenAIGPTImage1,
     "IdeogramTextToImage": IdeogramTextToImage,
+    "FluxProUltraImageNode": FluxProUltraImageNode,
     "MinimaxTextToVideoNode": MinimaxTextToVideoNode,
 }
 
@@ -781,5 +955,6 @@ 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",
 }