173 lines
6.7 KiB
Python
173 lines
6.7 KiB
Python
import aiohttp
|
||
import asyncio
|
||
import base64
|
||
import logging
|
||
from typing import Optional, Dict, Any
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
async def send_bfl_image_generation_request(
|
||
api_key: str,
|
||
prompt: str,
|
||
aspect_ratio: str = "1:1",
|
||
# Editing support
|
||
input_image_base64: Optional[str] = None,
|
||
# Internal tuning
|
||
seed: Optional[int] = None,
|
||
prompt_upsampling: bool = False,
|
||
safety_tolerance: int = 2,
|
||
output_format: Optional[str] = None,
|
||
polling_interval: float = 0.5,
|
||
max_wait_seconds: int = 120,
|
||
) -> Dict[str, Any]:
|
||
"""Submit a text-to-image or image-edit request to the BFL FLUX Kontext MAX endpoint.
|
||
|
||
The helper wraps the asynchronous flow required by the BFL API:
|
||
1. Send a POST request to ``/flux-kontext-max`` with the requested parameters.
|
||
2. Poll the provided ``polling_url`` until the generation is *Ready* (or fails).
|
||
3. Download the resulting image, convert it to **Base64** and re-shape the
|
||
response so that downstream utilities (e.g. *process_images_for_comfy*)
|
||
can handle it transparently. The returned structure mimics the OpenAI
|
||
``{'data': [{'b64_json': ...}]}`` schema.
|
||
|
||
Parameters
|
||
----------
|
||
api_key : str
|
||
The BFL account API key. Required.
|
||
prompt : str
|
||
A description of the desired image or the edit operation.
|
||
aspect_ratio : str, optional
|
||
Ratio in the format ``"W:H"`` (e.g. ``"1:1"``, ``"16:9"``). Defaults
|
||
to ``"1:1"`` which results in a 1024×1024 output.
|
||
input_image_base64 : str or None, optional
|
||
When provided the request is interpreted as an *edit* operation. The
|
||
value **must** be a Base64-encoded string (without the data-URL prefix).
|
||
seed, prompt_upsampling, safety_tolerance, output_format : various, optional
|
||
Direct mappings of the official BFL parameters. ``None`` omits the
|
||
field allowing the server-side default (currently *png*).
|
||
polling_interval : float, optional
|
||
Delay between polling attempts in seconds. Defaults to **0.5 s**.
|
||
max_wait_seconds : int, optional
|
||
Upper bound for the total time spent polling. An exception is raised
|
||
when the limit is exceeded (default: **120 s**).
|
||
|
||
Returns
|
||
-------
|
||
dict
|
||
A dictionary shaped like the OpenAI response (``{"data": [{"b64_json": …}]}``).
|
||
"""
|
||
|
||
if not api_key:
|
||
raise ValueError("BFL image generation requires a valid API key")
|
||
|
||
# ------------------------------------------------------------------
|
||
# 1. Submit generation request
|
||
# ------------------------------------------------------------------
|
||
endpoint = "https://api.bfl.ai/v1/flux-kontext-max"
|
||
headers = {
|
||
"accept": "application/json",
|
||
"x-key": api_key,
|
||
"Content-Type": "application/json",
|
||
}
|
||
|
||
payload: Dict[str, Any] = {
|
||
"prompt": prompt,
|
||
}
|
||
if aspect_ratio:
|
||
payload["aspect_ratio"] = aspect_ratio
|
||
if seed is not None:
|
||
payload["seed"] = seed
|
||
if prompt_upsampling:
|
||
payload["prompt_upsampling"] = prompt_upsampling
|
||
if safety_tolerance is not None:
|
||
payload["safety_tolerance"] = safety_tolerance
|
||
if output_format is not None:
|
||
payload["output_format"] = output_format
|
||
# Image editing
|
||
if input_image_base64:
|
||
payload["input_image"] = input_image_base64
|
||
|
||
logger.info(
|
||
"BFL API generation request → %s | aspect_ratio=%s | edit=%s",
|
||
endpoint,
|
||
aspect_ratio,
|
||
bool(input_image_base64),
|
||
)
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post(endpoint, headers=headers, json=payload) as resp:
|
||
if resp.status != 200:
|
||
text = await resp.text()
|
||
raise RuntimeError(
|
||
f"BFL API error: status {resp.status} – {text[:200]}…"
|
||
)
|
||
first_response = await resp.json()
|
||
|
||
request_id = first_response.get("id")
|
||
polling_url = first_response.get("polling_url")
|
||
if not polling_url or not request_id:
|
||
raise RuntimeError(
|
||
f"Unexpected BFL response: missing id/polling_url → {first_response}"
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# 2. Poll for completion
|
||
# ------------------------------------------------------------------
|
||
poll_headers = {"accept": "application/json", "x-key": api_key}
|
||
elapsed = 0.0
|
||
status = None
|
||
result_json: Optional[Dict[str, Any]] = None
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
while elapsed < max_wait_seconds:
|
||
await asyncio.sleep(polling_interval)
|
||
elapsed += polling_interval
|
||
|
||
async with session.get(
|
||
polling_url, headers=poll_headers, params={"id": request_id}
|
||
) as poll_resp:
|
||
if poll_resp.status != 200:
|
||
text = await poll_resp.text()
|
||
raise RuntimeError(
|
||
f"BFL polling error: status {poll_resp.status} – {text[:200]}…"
|
||
)
|
||
result_json = await poll_resp.json()
|
||
status = result_json.get("status")
|
||
|
||
if status == "Ready":
|
||
break
|
||
if status in {"Error", "Failed"}:
|
||
raise RuntimeError(
|
||
f"BFL generation failed: {result_json}" # type: ignore[arg-type]
|
||
)
|
||
else:
|
||
raise TimeoutError(
|
||
f"BFL generation polling timed-out after {max_wait_seconds} seconds"
|
||
)
|
||
|
||
assert result_json is not None # mypy appeasement
|
||
|
||
sample_url: Optional[str] = (
|
||
result_json.get("result", {}).get("sample") if isinstance(result_json, dict) else None
|
||
)
|
||
if not sample_url:
|
||
raise RuntimeError(
|
||
f"BFL response missing sample URL → {result_json}" # type: ignore[arg-type]
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# 3. Download image & convert to Base64
|
||
# ------------------------------------------------------------------
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.get(sample_url) as img_resp:
|
||
if img_resp.status != 200:
|
||
raise RuntimeError(
|
||
f"Could not download generated image: HTTP {img_resp.status}"
|
||
)
|
||
img_bytes = await img_resp.read()
|
||
|
||
img_b64 = base64.b64encode(img_bytes).decode("ascii")
|
||
|
||
# Shape to OpenAI-like structure so downstream code can reuse the same logic
|
||
return {"data": [{"b64_json": img_b64}]} |