converted BFL API nodes to async

This commit is contained in:
bigcat88 2025-08-01 09:25:24 +03:00
parent 99363dc394
commit abfd72f862
No known key found for this signature in database
GPG Key ID: 1F0BF0EC3CF22721
2 changed files with 70 additions and 69 deletions

View File

@ -22,7 +22,6 @@ from server import PromptServer
import numpy as np import numpy as np
from PIL import Image from PIL import Image
import requests
import torch import torch
import math import math
import base64 import base64
@ -208,9 +207,9 @@ async def download_url_to_image_tensor(url: str, timeout: int = None) -> torch.T
return bytesio_to_image_tensor(image_bytesio) return bytesio_to_image_tensor(image_bytesio)
def process_image_response(response: requests.Response) -> torch.Tensor: def process_image_response(response_content: bytes | str) -> torch.Tensor:
"""Uses content from a Response object and converts it to a torch.Tensor""" """Uses content from a Response object and converts it to a torch.Tensor"""
return bytesio_to_image_tensor(BytesIO(response.content)) return bytesio_to_image_tensor(BytesIO(response_content))
def _tensor_to_pil(image: torch.Tensor, total_pixels: int = 2048 * 2048) -> Image.Image: def _tensor_to_pil(image: torch.Tensor, total_pixels: int = 2048 * 2048) -> Image.Image:

View File

@ -1,3 +1,4 @@
import asyncio
import io import io
from inspect import cleandoc from inspect import cleandoc
from typing import Union, Optional from typing import Union, Optional
@ -28,7 +29,7 @@ from comfy_api_nodes.apinode_utils import (
import numpy as np import numpy as np
from PIL import Image from PIL import Image
import requests import aiohttp
import torch import torch
import base64 import base64
import time import time
@ -44,18 +45,18 @@ def convert_mask_to_image(mask: torch.Tensor):
return mask return mask
def handle_bfl_synchronous_operation( async def handle_bfl_synchronous_operation(
operation: SynchronousOperation, operation: SynchronousOperation,
timeout_bfl_calls=360, timeout_bfl_calls=360,
node_id: Union[str, None] = None, node_id: Union[str, None] = None,
): ):
response_api: BFLFluxProGenerateResponse = operation.execute() response_api: BFLFluxProGenerateResponse = await operation.execute()
return _poll_until_generated( return await _poll_until_generated(
response_api.polling_url, timeout=timeout_bfl_calls, node_id=node_id response_api.polling_url, timeout=timeout_bfl_calls, node_id=node_id
) )
def _poll_until_generated( async def _poll_until_generated(
polling_url: str, timeout=360, node_id: Union[str, None] = None polling_url: str, timeout=360, node_id: Union[str, None] = None
): ):
# used bfl-comfy-nodes to verify code implementation: # used bfl-comfy-nodes to verify code implementation:
@ -66,55 +67,56 @@ def _poll_until_generated(
retry_404_seconds = 2 retry_404_seconds = 2
retry_202_seconds = 2 retry_202_seconds = 2
retry_pending_seconds = 1 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:
if node_id:
time_elapsed = time.time() - start_time
PromptServer.instance.send_progress_text(
f"Generating ({time_elapsed:.0f}s)", node_id
)
response = requests.Session().send(request.prepare()) async with aiohttp.ClientSession() as session:
if response.status_code == 200: # NOTE: should True loop be replaced with checking if workflow has been interrupted?
result = response.json() while True:
if result["status"] == BFLStatus.ready: if node_id:
img_url = result["result"]["sample"] time_elapsed = time.time() - start_time
if node_id: PromptServer.instance.send_progress_text(
PromptServer.instance.send_progress_text( f"Generating ({time_elapsed:.0f}s)", node_id
f"Result URL: {img_url}", node_id
)
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}.") async with session.get(polling_url) as response:
elif result["status"] == BFLStatus.pending: if response.status == 200:
time.sleep(retry_pending_seconds) result = await response.json()
continue if result["status"] == BFLStatus.ready:
elif response.status_code == 404: img_url = result["result"]["sample"]
if retries_404 < max_retries_404: if node_id:
retries_404 += 1 PromptServer.instance.send_progress_text(
time.sleep(retry_404_seconds) f"Result URL: {img_url}", node_id
continue )
raise Exception( async with session.get(img_url) as img_resp:
f"BFL API could not find task after {max_retries_404} tries." return process_image_response(await img_resp.content.read())
) elif result["status"] in [
elif response.status_code == 202: BFLStatus.request_moderated,
time.sleep(retry_202_seconds) BFLStatus.content_moderated,
elif time.time() - start_time > timeout: ]:
raise Exception( status = result["status"]
f"BFL API experienced a timeout; could not return request under {timeout} seconds." raise Exception(
) f"BFL API did not return an image due to: {status}."
else: )
raise Exception(f"BFL API encountered an error: {response.json()}") elif result["status"] == BFLStatus.error:
raise Exception(f"BFL API encountered an error: {result}.")
elif result["status"] == BFLStatus.pending:
await asyncio.sleep(retry_pending_seconds)
continue
elif response.status == 404:
if retries_404 < max_retries_404:
retries_404 += 1
await asyncio.sleep(retry_404_seconds)
continue
raise Exception(
f"BFL API could not find task after {max_retries_404} tries."
)
elif response.status == 202:
await asyncio.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(image: torch.Tensor): def convert_image_to_base64(image: torch.Tensor):
scaled_image = downscale_image_tensor(image, total_pixels=2048 * 2048) scaled_image = downscale_image_tensor(image, total_pixels=2048 * 2048)
@ -222,7 +224,7 @@ class FluxProUltraImageNode(ComfyNodeABC):
API_NODE = True API_NODE = True
CATEGORY = "api node/image/BFL" CATEGORY = "api node/image/BFL"
def api_call( async def api_call(
self, self,
prompt: str, prompt: str,
aspect_ratio: str, aspect_ratio: str,
@ -266,7 +268,7 @@ class FluxProUltraImageNode(ComfyNodeABC):
), ),
auth_kwargs=kwargs, auth_kwargs=kwargs,
) )
output_image = handle_bfl_synchronous_operation(operation, node_id=unique_id) output_image = await handle_bfl_synchronous_operation(operation, node_id=unique_id)
return (output_image,) return (output_image,)
@ -354,7 +356,7 @@ class FluxKontextProImageNode(ComfyNodeABC):
BFL_PATH = "/proxy/bfl/flux-kontext-pro/generate" BFL_PATH = "/proxy/bfl/flux-kontext-pro/generate"
def api_call( async def api_call(
self, self,
prompt: str, prompt: str,
aspect_ratio: str, aspect_ratio: str,
@ -397,7 +399,7 @@ class FluxKontextProImageNode(ComfyNodeABC):
), ),
auth_kwargs=kwargs, auth_kwargs=kwargs,
) )
output_image = handle_bfl_synchronous_operation(operation, node_id=unique_id) output_image = await handle_bfl_synchronous_operation(operation, node_id=unique_id)
return (output_image,) return (output_image,)
@ -489,7 +491,7 @@ class FluxProImageNode(ComfyNodeABC):
API_NODE = True API_NODE = True
CATEGORY = "api node/image/BFL" CATEGORY = "api node/image/BFL"
def api_call( async def api_call(
self, self,
prompt: str, prompt: str,
prompt_upsampling, prompt_upsampling,
@ -524,7 +526,7 @@ class FluxProImageNode(ComfyNodeABC):
), ),
auth_kwargs=kwargs, auth_kwargs=kwargs,
) )
output_image = handle_bfl_synchronous_operation(operation, node_id=unique_id) output_image = await handle_bfl_synchronous_operation(operation, node_id=unique_id)
return (output_image,) return (output_image,)
@ -632,7 +634,7 @@ class FluxProExpandNode(ComfyNodeABC):
API_NODE = True API_NODE = True
CATEGORY = "api node/image/BFL" CATEGORY = "api node/image/BFL"
def api_call( async def api_call(
self, self,
image: torch.Tensor, image: torch.Tensor,
prompt: str, prompt: str,
@ -670,7 +672,7 @@ class FluxProExpandNode(ComfyNodeABC):
), ),
auth_kwargs=kwargs, auth_kwargs=kwargs,
) )
output_image = handle_bfl_synchronous_operation(operation, node_id=unique_id) output_image = await handle_bfl_synchronous_operation(operation, node_id=unique_id)
return (output_image,) return (output_image,)
@ -744,7 +746,7 @@ class FluxProFillNode(ComfyNodeABC):
API_NODE = True API_NODE = True
CATEGORY = "api node/image/BFL" CATEGORY = "api node/image/BFL"
def api_call( async def api_call(
self, self,
image: torch.Tensor, image: torch.Tensor,
mask: torch.Tensor, mask: torch.Tensor,
@ -780,7 +782,7 @@ class FluxProFillNode(ComfyNodeABC):
), ),
auth_kwargs=kwargs, auth_kwargs=kwargs,
) )
output_image = handle_bfl_synchronous_operation(operation, node_id=unique_id) output_image = await handle_bfl_synchronous_operation(operation, node_id=unique_id)
return (output_image,) return (output_image,)
@ -879,7 +881,7 @@ class FluxProCannyNode(ComfyNodeABC):
API_NODE = True API_NODE = True
CATEGORY = "api node/image/BFL" CATEGORY = "api node/image/BFL"
def api_call( async def api_call(
self, self,
control_image: torch.Tensor, control_image: torch.Tensor,
prompt: str, prompt: str,
@ -929,7 +931,7 @@ class FluxProCannyNode(ComfyNodeABC):
), ),
auth_kwargs=kwargs, auth_kwargs=kwargs,
) )
output_image = handle_bfl_synchronous_operation(operation, node_id=unique_id) output_image = await handle_bfl_synchronous_operation(operation, node_id=unique_id)
return (output_image,) return (output_image,)
@ -1008,7 +1010,7 @@ class FluxProDepthNode(ComfyNodeABC):
API_NODE = True API_NODE = True
CATEGORY = "api node/image/BFL" CATEGORY = "api node/image/BFL"
def api_call( async def api_call(
self, self,
control_image: torch.Tensor, control_image: torch.Tensor,
prompt: str, prompt: str,
@ -1045,7 +1047,7 @@ class FluxProDepthNode(ComfyNodeABC):
), ),
auth_kwargs=kwargs, auth_kwargs=kwargs,
) )
output_image = handle_bfl_synchronous_operation(operation, node_id=unique_id) output_image = await handle_bfl_synchronous_operation(operation, node_id=unique_id)
return (output_image,) return (output_image,)