mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
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:
parent
8665e92aa8
commit
cd731811d9
9 changed files with 832 additions and 0 deletions
19
litellm/llms/black_forest_labs/__init__.py
Normal file
19
litellm/llms/black_forest_labs/__init__.py
Normal 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",
|
||||
]
|
||||
39
litellm/llms/black_forest_labs/common_utils.py
Normal file
39
litellm/llms/black_forest_labs/common_utils.py
Normal 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",
|
||||
}
|
||||
3
litellm/llms/black_forest_labs/image_edit/__init__.py
Normal file
3
litellm/llms/black_forest_labs/image_edit/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
from .transformation import BlackForestLabsImageEditConfig
|
||||
|
||||
__all__ = ["BlackForestLabsImageEditConfig"]
|
||||
334
litellm/llms/black_forest_labs/image_edit/transformation.py
Normal file
334
litellm/llms/black_forest_labs/image_edit/transformation.py
Normal 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)],
|
||||
)
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
0
tests/test_litellm/llms/black_forest_labs/__init__.py
Normal file
0
tests/test_litellm/llms/black_forest_labs/__init__.py
Normal 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
|
||||
Loading…
Add table
Reference in a new issue