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,7 +67,8 @@ 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)
async with aiohttp.ClientSession() as session:
# NOTE: should True loop be replaced with checking if workflow has been interrupted? # NOTE: should True loop be replaced with checking if workflow has been interrupted?
while True: while True:
if node_id: if node_id:
@ -75,17 +77,17 @@ def _poll_until_generated(
f"Generating ({time_elapsed:.0f}s)", node_id f"Generating ({time_elapsed:.0f}s)", node_id
) )
response = requests.Session().send(request.prepare()) async with session.get(polling_url) as response:
if response.status_code == 200: if response.status == 200:
result = response.json() result = await response.json()
if result["status"] == BFLStatus.ready: if result["status"] == BFLStatus.ready:
img_url = result["result"]["sample"] img_url = result["result"]["sample"]
if node_id: if node_id:
PromptServer.instance.send_progress_text( PromptServer.instance.send_progress_text(
f"Result URL: {img_url}", node_id f"Result URL: {img_url}", node_id
) )
img_response = requests.get(img_url) async with session.get(img_url) as img_resp:
return process_image_response(img_response) return process_image_response(await img_resp.content.read())
elif result["status"] in [ elif result["status"] in [
BFLStatus.request_moderated, BFLStatus.request_moderated,
BFLStatus.content_moderated, BFLStatus.content_moderated,
@ -97,18 +99,18 @@ def _poll_until_generated(
elif result["status"] == BFLStatus.error: elif result["status"] == BFLStatus.error:
raise Exception(f"BFL API encountered an error: {result}.") raise Exception(f"BFL API encountered an error: {result}.")
elif result["status"] == BFLStatus.pending: elif result["status"] == BFLStatus.pending:
time.sleep(retry_pending_seconds) await asyncio.sleep(retry_pending_seconds)
continue continue
elif response.status_code == 404: elif response.status == 404:
if retries_404 < max_retries_404: if retries_404 < max_retries_404:
retries_404 += 1 retries_404 += 1
time.sleep(retry_404_seconds) await asyncio.sleep(retry_404_seconds)
continue continue
raise Exception( raise Exception(
f"BFL API could not find task after {max_retries_404} tries." f"BFL API could not find task after {max_retries_404} tries."
) )
elif response.status_code == 202: elif response.status == 202:
time.sleep(retry_202_seconds) await asyncio.sleep(retry_202_seconds)
elif time.time() - start_time > timeout: elif time.time() - start_time > timeout:
raise Exception( raise Exception(
f"BFL API experienced a timeout; could not return request under {timeout} seconds." f"BFL API experienced a timeout; could not return request under {timeout} seconds."
@ -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,)