123 lines
4.3 KiB
Python
123 lines
4.3 KiB
Python
import time
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import requests
|
|
|
|
from .flux2_config import Flux2Config
|
|
|
|
|
|
class Flux2APIError(Exception):
|
|
"""Raised when the FLUX.2 API returns an error or an unexpected response."""
|
|
|
|
|
|
class Flux2API:
|
|
"""Minimal client for the Black Forest Labs FLUX.2 [pro] endpoint."""
|
|
|
|
BASE_URL = "https://api.bfl.ai/v1/flux-2-pro"
|
|
|
|
def __init__(self, timeout: int = 600, poll_interval: float = 1.0):
|
|
self.timeout = timeout
|
|
self.poll_interval = poll_interval
|
|
|
|
api_key = Flux2Config().get_key()
|
|
if not api_key:
|
|
raise Flux2APIError(
|
|
"BFL_API_KEY is not set. Define it in config.ini or export BFL_API_KEY before running ComfyUI."
|
|
)
|
|
|
|
self.headers = {
|
|
"accept": "application/json",
|
|
"Content-Type": "application/json",
|
|
"x-key": api_key,
|
|
}
|
|
|
|
@staticmethod
|
|
def _strip_empty(payload: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Remove keys with None or empty-string values so we only send what the API needs."""
|
|
cleaned: Dict[str, Any] = {}
|
|
for key, value in payload.items():
|
|
if value is None:
|
|
continue
|
|
if isinstance(value, str) and value.strip() == "":
|
|
continue
|
|
cleaned[key] = value
|
|
return cleaned
|
|
|
|
def submit_request(self, payload: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Submit a generation/editing request and return the raw API response."""
|
|
body = self._strip_empty(payload)
|
|
response = requests.post(self.BASE_URL, headers=self.headers, json=body, timeout=30)
|
|
response.raise_for_status()
|
|
|
|
data = response.json()
|
|
if "polling_url" not in data:
|
|
raise Flux2APIError("FLUX.2 API response did not include a polling_url.")
|
|
if "id" not in data:
|
|
raise Flux2APIError("FLUX.2 API response did not include a request id.")
|
|
return data
|
|
|
|
def poll_result(self, polling_url: str, request_id: Optional[str] = None) -> Dict[str, Any]:
|
|
"""Poll the provided polling_url until the job is ready or fails."""
|
|
start_time = time.time()
|
|
last_status: Optional[str] = None
|
|
|
|
while (time.time() - start_time) < self.timeout:
|
|
response = requests.get(
|
|
polling_url,
|
|
headers=self.headers,
|
|
params={"id": request_id} if request_id else None,
|
|
timeout=30,
|
|
)
|
|
response.raise_for_status()
|
|
payload = response.json()
|
|
|
|
status = payload.get("status")
|
|
if status != last_status:
|
|
print(f"[Flux2] Status: {status}")
|
|
last_status = status
|
|
|
|
if status == "Ready":
|
|
return payload
|
|
if status in {"Error", "Failed", "Content Moderated", "Request Moderated"}:
|
|
raise Flux2APIError(f"FLUX.2 request failed with status '{status}': {payload}")
|
|
|
|
time.sleep(self.poll_interval)
|
|
|
|
raise Flux2APIError(f"Timed out after {self.timeout}s while waiting for the FLUX.2 result.")
|
|
|
|
def run(self, payload: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Submit a request and wait for completion, returning convenience fields and raw payloads."""
|
|
initial = self.submit_request(payload)
|
|
polling_url = initial["polling_url"]
|
|
request_id = initial.get("id")
|
|
print(f"[Flux2] Submitted request: {request_id}")
|
|
|
|
final = self.poll_result(polling_url, request_id=request_id)
|
|
sample_url = final.get("result", {}).get("sample")
|
|
cost = initial.get("cost")
|
|
return {
|
|
"sample": sample_url,
|
|
"cost": cost,
|
|
"request": initial,
|
|
"result": final,
|
|
}
|
|
|
|
|
|
def merge_reference_images(main_image: str, additional_images: List[str]) -> Dict[str, str]:
|
|
"""
|
|
Build the payload keys for multi-reference inputs.
|
|
|
|
FLUX.2 [pro] accepts input_image + up to 7 extras: input_image_2 ... input_image_8.
|
|
"""
|
|
payload: Dict[str, str] = {}
|
|
if main_image:
|
|
payload["input_image"] = main_image
|
|
|
|
for idx, image in enumerate(additional_images, start=2):
|
|
if idx > 8:
|
|
break
|
|
if image and image.strip():
|
|
payload[f"input_image_{idx}"] = image.strip()
|
|
|
|
return payload
|