feat(black_forest_labs): add native image edit support for Black Forest Labs

Add native integration for Black Forest Labs image editing models
(flux-kontext-pro, flux-kontext-max, flux-pro-1.0-fill, flux-pro-1.0-expand).

Changes:
- Add BlackForestLabsImageEditConfig for BFL API transformation
- Add BLACK_FOREST_LABS to LlmProviders enum
- Add use_multipart_form_data() to BaseImageEditConfig for JSON vs form-data
- Modify image_edit_handler to support JSON request bodies
- Add comprehensive unit tests

Closes #11401
This commit is contained in:
Chesars 2025-12-15 17:50:12 -03:00
parent 8665e92aa8
commit cd731811d9
9 changed files with 832 additions and 0 deletions

View file

@ -0,0 +1,19 @@
from .common_utils import (
DEFAULT_API_BASE,
DEFAULT_MAX_POLLING_TIME,
DEFAULT_POLLING_INTERVAL,
IMAGE_EDIT_MODELS,
IMAGE_GENERATION_MODELS,
BlackForestLabsError,
)
from .image_edit import BlackForestLabsImageEditConfig
__all__ = [
"BlackForestLabsError",
"BlackForestLabsImageEditConfig",
"DEFAULT_API_BASE",
"DEFAULT_MAX_POLLING_TIME",
"DEFAULT_POLLING_INTERVAL",
"IMAGE_EDIT_MODELS",
"IMAGE_GENERATION_MODELS",
]

View file

@ -0,0 +1,39 @@
"""
Black Forest Labs Common Utilities
Common utilities, constants, and error handling for Black Forest Labs API.
"""
from typing import Dict
from litellm.llms.base_llm.chat.transformation import BaseLLMException
class BlackForestLabsError(BaseLLMException):
"""Exception class for Black Forest Labs API errors."""
pass
# API Constants
DEFAULT_API_BASE = "https://api.bfl.ai"
# Polling configuration
DEFAULT_POLLING_INTERVAL = 1.5 # seconds
DEFAULT_MAX_POLLING_TIME = 300 # 5 minutes
# Model to endpoint mapping for image edit
IMAGE_EDIT_MODELS: Dict[str, str] = {
"flux-kontext-pro": "/v1/flux-kontext-pro",
"flux-kontext-max": "/v1/flux-kontext-max",
"flux-pro-1.0-fill": "/v1/flux-pro-1.0-fill",
"flux-pro-1.0-expand": "/v1/flux-pro-1.0-expand",
}
# Model to endpoint mapping for image generation
IMAGE_GENERATION_MODELS: Dict[str, str] = {
"flux-pro-1.1": "/v1/flux-pro-1.1",
"flux-pro-1.1-ultra": "/v1/flux-pro-1.1-ultra",
"flux-dev": "/v1/flux-dev",
"flux-pro": "/v1/flux-pro",
}

View file

@ -0,0 +1,3 @@
from .transformation import BlackForestLabsImageEditConfig
__all__ = ["BlackForestLabsImageEditConfig"]

View file

@ -0,0 +1,334 @@
"""
Black Forest Labs Image Edit Configuration
Handles transformation between OpenAI-compatible format and Black Forest Labs API format
for image editing endpoints (flux-kontext-pro, flux-kontext-max, etc.).
API Reference: https://docs.bfl.ai/
"""
import base64
import time
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
import httpx
from httpx._types import RequestFiles
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import FileTypes, ImageObject, ImageResponse
from ..common_utils import (
DEFAULT_API_BASE,
DEFAULT_MAX_POLLING_TIME,
DEFAULT_POLLING_INTERVAL,
IMAGE_EDIT_MODELS,
BlackForestLabsError,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class BlackForestLabsImageEditConfig(BaseImageEditConfig):
"""
Configuration for Black Forest Labs image editing.
Supports:
- flux-kontext-pro: General image editing with prompts
- flux-kontext-max: Premium quality editing
- flux-pro-1.0-fill: Inpainting with mask
- flux-pro-1.0-expand: Outpainting (expand image borders)
"""
def get_supported_openai_params(self, model: str) -> List[str]:
"""
Return list of OpenAI params supported by Black Forest Labs.
Note: BFL uses different parameter names, these are mapped in map_openai_params.
"""
return [
"n", # Number of images (BFL returns 1 per request)
"size", # Maps to aspect_ratio
"response_format", # b64_json or url
]
def map_openai_params(
self,
image_edit_optional_params: ImageEditOptionalRequestParams,
model: str,
drop_params: bool,
) -> Dict:
"""
Map OpenAI parameters to Black Forest Labs parameters.
BFL-specific params are passed through directly.
"""
optional_params: Dict[str, Any] = {}
# Pass through BFL-specific params
bfl_params = [
"seed",
"output_format",
"safety_tolerance",
"prompt_upsampling",
# Kontext-specific
"aspect_ratio",
# Fill/Inpaint-specific
"steps",
"guidance",
"grow_mask",
# Expand-specific
"top",
"bottom",
"left",
"right",
]
# Convert TypedDict to regular dict for access
params_dict = dict(image_edit_optional_params)
for param in bfl_params:
if param in params_dict:
value = params_dict[param]
if value is not None:
optional_params[param] = value
# Set default output format
if "output_format" not in optional_params:
optional_params["output_format"] = "png"
return optional_params
def validate_environment(
self,
headers: dict,
model: str,
api_key: Optional[str] = None,
) -> dict:
"""
Validate environment and set up headers for Black Forest Labs.
BFL uses x-key header for authentication.
"""
final_api_key: Optional[str] = (
api_key
or get_secret_str("BFL_API_KEY")
or get_secret_str("BLACK_FOREST_LABS_API_KEY")
)
if not final_api_key:
raise BlackForestLabsError(
status_code=401,
message="BFL_API_KEY is not set. Please set it via environment variable or pass api_key parameter.",
)
headers["x-key"] = final_api_key
headers["Content-Type"] = "application/json"
headers["Accept"] = "application/json"
return headers
def use_multipart_form_data(self) -> bool:
"""
BFL uses JSON requests, not multipart/form-data.
"""
return False
def _get_model_endpoint(self, model: str) -> str:
"""
Get the API endpoint for a given model.
"""
# Remove provider prefix if present (e.g., "black_forest_labs/flux-kontext-pro")
model_name = model.lower()
if "/" in model_name:
model_name = model_name.split("/")[-1]
# Check if model is in our mapping
if model_name in IMAGE_EDIT_MODELS:
return IMAGE_EDIT_MODELS[model_name]
# Default to kontext-pro
return IMAGE_EDIT_MODELS["flux-kontext-pro"]
def get_complete_url(
self,
model: str,
api_base: Optional[str],
litellm_params: dict,
) -> str:
"""
Get the complete URL for the Black Forest Labs API request.
"""
base_url: str = (
api_base
or get_secret_str("BFL_API_BASE")
or DEFAULT_API_BASE
)
base_url = base_url.rstrip("/")
endpoint = self._get_model_endpoint(model)
return f"{base_url}{endpoint}"
def _read_image_bytes(self, image: Any) -> bytes:
"""Read image bytes from various input types."""
if isinstance(image, bytes):
return image
elif isinstance(image, list):
# If it's a list, take the first image
return self._read_image_bytes(image[0])
elif hasattr(image, "read"):
# File-like object
pos = getattr(image, "tell", lambda: 0)()
if hasattr(image, "seek"):
image.seek(0)
data = image.read()
if hasattr(image, "seek"):
image.seek(pos)
return data
else:
return image
def transform_image_edit_request(
self,
model: str,
prompt: str,
image: FileTypes,
image_edit_optional_request_params: Dict,
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> Tuple[Dict, RequestFiles]:
"""
Transform OpenAI-style request to Black Forest Labs request format.
BFL uses JSON body with base64-encoded images, not multipart/form-data.
"""
# Read and encode image
image_bytes = self._read_image_bytes(image)
b64_image = base64.b64encode(image_bytes).decode("utf-8")
# Build request body
request_body: Dict[str, Any] = {
"prompt": prompt,
"input_image": b64_image,
}
# Add optional params
for key, value in image_edit_optional_request_params.items():
if key not in ["extra_headers", "extra_body"] and value is not None:
request_body[key] = value
# Handle mask if provided (for inpainting)
if "mask" in image_edit_optional_request_params:
mask = image_edit_optional_request_params["mask"]
mask_bytes = self._read_image_bytes(mask)
request_body["mask"] = base64.b64encode(mask_bytes).decode("utf-8")
# BFL uses JSON, not multipart - return empty files
return request_body, []
def _poll_for_result(
self,
polling_url: str,
api_key: str,
max_wait: float = DEFAULT_MAX_POLLING_TIME,
interval: float = DEFAULT_POLLING_INTERVAL,
) -> Dict:
"""
Poll the BFL API until the result is ready.
Returns the result data when status is "Ready".
Raises BlackForestLabsError on failure.
"""
start_time = time.time()
while time.time() - start_time < max_wait:
response = httpx.get(
polling_url,
headers={"x-key": api_key},
timeout=30.0,
)
if response.status_code != 200:
raise BlackForestLabsError(
status_code=response.status_code,
message=f"Polling failed: {response.text}",
)
data = response.json()
status = data.get("status")
if status == "Ready":
return data
elif status in ["Error", "Content Moderated", "Request Moderated"]:
raise BlackForestLabsError(
status_code=400,
message=f"Image generation failed: {status}",
)
time.sleep(interval)
raise BlackForestLabsError(
status_code=408,
message=f"Timeout waiting for result after {max_wait} seconds",
)
def transform_image_edit_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
) -> ImageResponse:
"""
Transform Black Forest Labs response to OpenAI-compatible ImageResponse.
BFL returns a task ID initially, then we poll until the result is ready.
"""
try:
response_data = raw_response.json()
except Exception as e:
raise BlackForestLabsError(
status_code=raw_response.status_code,
message=f"Error parsing BFL response: {e}",
)
# Check for immediate errors
if "errors" in response_data:
raise BlackForestLabsError(
status_code=raw_response.status_code,
message=f"BFL error: {response_data['errors']}",
)
# Get polling URL
polling_url = response_data.get("polling_url")
if not polling_url:
raise BlackForestLabsError(
status_code=500,
message="No polling_url in BFL response",
)
# Extract API key from original request headers
api_key = raw_response.request.headers.get("x-key", "")
# Poll for result
result_data = self._poll_for_result(polling_url, api_key)
# Get image URL from result
image_url = result_data.get("result", {}).get("sample")
if not image_url:
raise BlackForestLabsError(
status_code=500,
message="No image URL in BFL result",
)
# Build ImageResponse
return ImageResponse(
created=int(time.time()),
data=[ImageObject(url=image_url)],
)

View file

@ -3099,6 +3099,7 @@ class LlmProviders(str, Enum):
GEMINI = "gemini"
AI21 = "ai21"
BASETEN = "baseten"
BLACK_FOREST_LABS = "black_forest_labs"
AZURE = "azure"
AZURE_TEXT = "azure_text"
AZURE_AI = "azure_ai"

View file

@ -8748,6 +8748,12 @@ class ProviderConfigManager:
)
return RecraftImageEditConfig()
elif LlmProviders.BLACK_FOREST_LABS == provider:
from litellm.llms.black_forest_labs.image_edit.transformation import (
BlackForestLabsImageEditConfig,
)
return BlackForestLabsImageEditConfig()
elif LlmProviders.AZURE_AI == provider:
from litellm.llms.azure_ai.image_edit import get_azure_ai_image_edit_config

View file

@ -0,0 +1,430 @@
"""
Unit tests for Black Forest Labs image edit transformation functionality.
"""
import base64
import json
import os
import sys
import time
from io import BytesIO
from typing import Dict, List
from unittest.mock import MagicMock, patch
import httpx
import pytest
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
from litellm.llms.black_forest_labs.image_edit.transformation import (
BlackForestLabsImageEditConfig,
)
from litellm.llms.black_forest_labs.common_utils import BlackForestLabsError
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ImageObject, ImageResponse
class TestBlackForestLabsImageEditTransformation:
"""
Unit tests for Black Forest Labs image edit transformation functionality.
"""
def setup_method(self):
"""Set up test fixtures before each test method."""
self.config = BlackForestLabsImageEditConfig()
self.model = "flux-kontext-pro"
self.logging_obj = MagicMock()
self.prompt = "Add a red hat to the person in the image"
def test_get_supported_openai_params(self):
"""Test that supported OpenAI params are returned correctly."""
params = self.config.get_supported_openai_params(self.model)
assert "n" in params
assert "size" in params
assert "response_format" in params
def test_map_openai_params_basic(self):
"""Test mapping of OpenAI params to BFL params."""
optional_params = ImageEditOptionalRequestParams()
result = self.config.map_openai_params(
image_edit_optional_params=optional_params,
model=self.model,
drop_params=False,
)
# Should have default output_format
assert result.get("output_format") == "png"
def test_map_openai_params_with_bfl_specific(self):
"""Test that BFL-specific params are passed through."""
# BFL-specific params are passed as dict keys
optional_params: ImageEditOptionalRequestParams = {
"seed": 42,
"safety_tolerance": 2,
"aspect_ratio": "16:9",
}
result = self.config.map_openai_params(
image_edit_optional_params=optional_params,
model=self.model,
drop_params=False,
)
assert result.get("seed") == 42
assert result.get("safety_tolerance") == 2
assert result.get("aspect_ratio") == "16:9"
assert result.get("output_format") == "png"
def test_validate_environment_with_api_key(self):
"""Test environment validation with provided API key."""
headers = {}
result = self.config.validate_environment(
headers=headers,
model=self.model,
api_key="test-api-key",
)
assert result["x-key"] == "test-api-key"
assert result["Content-Type"] == "application/json"
assert result["Accept"] == "application/json"
def test_validate_environment_missing_api_key(self):
"""Test that missing API key raises error."""
headers = {}
with patch("litellm.llms.black_forest_labs.image_edit.transformation.get_secret_str") as mock_get_secret:
mock_get_secret.return_value = None
with pytest.raises(BlackForestLabsError) as exc_info:
self.config.validate_environment(
headers=headers,
model=self.model,
api_key=None,
)
assert exc_info.value.status_code == 401
assert "BFL_API_KEY is not set" in exc_info.value.message
def test_get_model_endpoint_kontext_pro(self):
"""Test endpoint resolution for flux-kontext-pro."""
endpoint = self.config._get_model_endpoint("flux-kontext-pro")
assert endpoint == "/v1/flux-kontext-pro"
def test_get_model_endpoint_kontext_max(self):
"""Test endpoint resolution for flux-kontext-max."""
endpoint = self.config._get_model_endpoint("flux-kontext-max")
assert endpoint == "/v1/flux-kontext-max"
def test_get_model_endpoint_with_provider_prefix(self):
"""Test endpoint resolution with provider prefix."""
endpoint = self.config._get_model_endpoint("black_forest_labs/flux-kontext-pro")
assert endpoint == "/v1/flux-kontext-pro"
def test_get_model_endpoint_fill(self):
"""Test endpoint resolution for flux-pro-1.0-fill."""
endpoint = self.config._get_model_endpoint("flux-pro-1.0-fill")
assert endpoint == "/v1/flux-pro-1.0-fill"
def test_get_complete_url(self):
"""Test complete URL generation."""
url = self.config.get_complete_url(
model="flux-kontext-pro",
api_base=None,
litellm_params={},
)
assert url == "https://api.bfl.ai/v1/flux-kontext-pro"
def test_get_complete_url_custom_base(self):
"""Test complete URL generation with custom base."""
url = self.config.get_complete_url(
model="flux-kontext-pro",
api_base="https://custom.api.com/",
litellm_params={},
)
assert url == "https://custom.api.com/v1/flux-kontext-pro"
def test_transform_image_edit_request(self):
"""Test request transformation to BFL format."""
image_data = b"fake_image_data"
image = BytesIO(image_data)
image_edit_optional_params = {
"seed": 123,
"output_format": "jpeg",
}
litellm_params = GenericLiteLLMParams()
headers = {}
data, files = self.config.transform_image_edit_request(
model=self.model,
prompt=self.prompt,
image=image,
image_edit_optional_request_params=image_edit_optional_params,
litellm_params=litellm_params,
headers=headers,
)
# Check that data contains the expected parameters
assert data["prompt"] == self.prompt
assert "input_image" in data
# Verify base64 encoding
decoded = base64.b64decode(data["input_image"])
assert decoded == image_data
assert data["seed"] == 123
assert data["output_format"] == "jpeg"
# BFL uses JSON, not multipart - files should be empty
assert files == []
def test_transform_image_edit_request_with_mask(self):
"""Test request transformation with mask for inpainting."""
image_data = b"fake_image_data"
mask_data = b"fake_mask_data"
image = BytesIO(image_data)
image_edit_optional_params = {
"mask": BytesIO(mask_data),
"output_format": "png",
}
litellm_params = GenericLiteLLMParams()
headers = {}
data, files = self.config.transform_image_edit_request(
model="flux-pro-1.0-fill",
prompt=self.prompt,
image=image,
image_edit_optional_request_params=image_edit_optional_params,
litellm_params=litellm_params,
headers=headers,
)
# Check mask is base64 encoded
assert "mask" in data
decoded_mask = base64.b64decode(data["mask"])
assert decoded_mask == mask_data
def test_read_image_bytes_from_bytes(self):
"""Test reading image bytes from bytes input."""
image_data = b"test_image_bytes"
result = self.config._read_image_bytes(image_data)
assert result == image_data
def test_read_image_bytes_from_file_like(self):
"""Test reading image bytes from file-like object."""
image_data = b"test_image_bytes"
image = BytesIO(image_data)
result = self.config._read_image_bytes(image)
assert result == image_data
def test_read_image_bytes_from_list(self):
"""Test reading image bytes from list (takes first)."""
image_data = b"test_image_bytes"
images = [BytesIO(image_data), BytesIO(b"other")]
result = self.config._read_image_bytes(images)
assert result == image_data
def test_poll_for_result_success(self):
"""Test successful polling."""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"status": "Ready",
"result": {"sample": "https://example.com/image.png"},
}
with patch("httpx.get", return_value=mock_response):
result = self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert result["status"] == "Ready"
assert result["result"]["sample"] == "https://example.com/image.png"
def test_poll_for_result_pending_then_ready(self):
"""Test polling that starts pending then becomes ready."""
pending_response = MagicMock()
pending_response.status_code = 200
pending_response.json.return_value = {"status": "Pending"}
ready_response = MagicMock()
ready_response.status_code = 200
ready_response.json.return_value = {
"status": "Ready",
"result": {"sample": "https://example.com/image.png"},
}
with patch("httpx.get", side_effect=[pending_response, ready_response]):
result = self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert result["status"] == "Ready"
def test_poll_for_result_error_status(self):
"""Test polling with error status."""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"status": "Error"}
with patch("httpx.get", return_value=mock_response):
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert exc_info.value.status_code == 400
assert "Error" in exc_info.value.message
def test_poll_for_result_content_moderated(self):
"""Test polling with content moderated status."""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"status": "Content Moderated"}
with patch("httpx.get", return_value=mock_response):
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert exc_info.value.status_code == 400
assert "Content Moderated" in exc_info.value.message
def test_poll_for_result_timeout(self):
"""Test polling timeout."""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"status": "Pending"}
with patch("httpx.get", return_value=mock_response):
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=0.2,
interval=0.1,
)
assert exc_info.value.status_code == 408
assert "Timeout" in exc_info.value.message
def test_poll_for_result_http_error(self):
"""Test polling with HTTP error."""
mock_response = MagicMock()
mock_response.status_code = 500
mock_response.text = "Internal Server Error"
with patch("httpx.get", return_value=mock_response):
with pytest.raises(BlackForestLabsError) as exc_info:
self.config._poll_for_result(
polling_url="https://api.bfl.ai/v1/get_result?id=123",
api_key="test-key",
max_wait=10,
interval=0.1,
)
assert exc_info.value.status_code == 500
def test_transform_image_edit_response_success(self):
"""Test successful response transformation."""
# Create mock initial response with polling URL
mock_request = MagicMock()
mock_request.headers = {"x-key": "test-key"}
mock_response = MagicMock()
mock_response.json.return_value = {
"id": "task-123",
"polling_url": "https://api.bfl.ai/v1/get_result?id=task-123",
}
mock_response.request = mock_request
mock_response.status_code = 200
# Mock the polling result
poll_response = MagicMock()
poll_response.status_code = 200
poll_response.json.return_value = {
"status": "Ready",
"result": {"sample": "https://example.com/edited-image.png"},
}
with patch("httpx.get", return_value=poll_response):
result = self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
assert isinstance(result, ImageResponse)
assert len(result.data) == 1
assert result.data[0].url == "https://example.com/edited-image.png"
assert result.created is not None
def test_transform_image_edit_response_no_polling_url(self):
"""Test response transformation when polling URL is missing."""
mock_response = MagicMock()
mock_response.json.return_value = {"id": "task-123"} # No polling_url
mock_response.status_code = 200
with pytest.raises(BlackForestLabsError) as exc_info:
self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
assert exc_info.value.status_code == 500
assert "No polling_url" in exc_info.value.message
def test_transform_image_edit_response_api_error(self):
"""Test response transformation with API error."""
mock_response = MagicMock()
mock_response.json.return_value = {
"errors": ["Invalid image format"]
}
mock_response.status_code = 400
with pytest.raises(BlackForestLabsError) as exc_info:
self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
assert "Invalid image format" in exc_info.value.message
def test_transform_image_edit_response_json_parse_error(self):
"""Test response transformation with JSON parse error."""
mock_response = MagicMock()
mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "", 0)
mock_response.status_code = 500
with pytest.raises(BlackForestLabsError) as exc_info:
self.config.transform_image_edit_response(
model=self.model,
raw_response=mock_response,
logging_obj=self.logging_obj,
)
assert "Error parsing BFL response" in exc_info.value.message