mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
Merge pull request #39424 from emerzon/litellm_azure_ai_flux_2_flex
feat(azure_ai): support FLUX.2 flex images
This commit is contained in:
Mateo Wang
GitHub
commit
a6bd779bd1
12 changed files with 574 additions and 58 deletions
|
|
@ -846,7 +846,12 @@ def image_edit(
|
|||
local_vars.update(kwargs)
|
||||
# Get ImageEditOptionalRequestParams with only valid parameters
|
||||
image_edit_optional_params: Final[ImageEditOptionalRequestParams] = (
|
||||
_get_ImageEditRequestUtils().get_requested_image_edit_optional_param(local_vars)
|
||||
_get_ImageEditRequestUtils().get_requested_image_edit_optional_param(
|
||||
local_vars,
|
||||
provider_supported_params=frozenset(
|
||||
image_edit_provider_config.get_supported_openai_params(model)
|
||||
).intersection(non_default_params),
|
||||
)
|
||||
)
|
||||
# Get optional parameters for the responses API
|
||||
image_edit_request_params: Final[dict] = _get_ImageEditRequestUtils().get_optional_params_image_edit(
|
||||
|
|
@ -857,7 +862,7 @@ def image_edit(
|
|||
additional_drop_params=kwargs.get("additional_drop_params"),
|
||||
)
|
||||
|
||||
if (
|
||||
if image_edit_provider_config.use_multipart_form_data() and (
|
||||
custom_llm_provider == "openai"
|
||||
or custom_llm_provider == "azure"
|
||||
or custom_llm_provider in litellm.openai_compatible_providers
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from collections.abc import Mapping
|
||||
from collections.abc import Collection, Mapping
|
||||
from io import BufferedReader, BytesIO
|
||||
from typing import Any, Final, cast, get_type_hints
|
||||
|
||||
|
|
@ -63,6 +63,7 @@ class ImageEditRequestUtils:
|
|||
@staticmethod
|
||||
def get_requested_image_edit_optional_param(
|
||||
params: Mapping[str, object],
|
||||
provider_supported_params: Collection[str] = (),
|
||||
) -> ImageEditOptionalRequestParams:
|
||||
"""
|
||||
Filter parameters to only include those defined in ImageEditOptionalRequestParams.
|
||||
|
|
@ -73,7 +74,9 @@ class ImageEditRequestUtils:
|
|||
Returns:
|
||||
ImageEditOptionalRequestParams instance with only the valid parameters
|
||||
"""
|
||||
valid_keys: Final = get_type_hints(ImageEditOptionalRequestParams).keys()
|
||||
valid_keys: Final = frozenset(get_type_hints(ImageEditOptionalRequestParams)) | frozenset(
|
||||
provider_supported_params
|
||||
)
|
||||
filtered_params: Final = {k: v for k, v in params.items() if k in valid_keys and v is not None}
|
||||
return cast(ImageEditOptionalRequestParams, filtered_params)
|
||||
|
||||
|
|
|
|||
|
|
@ -1848,6 +1848,9 @@ class CostCalculatorUtils:
|
|||
return azure_ai_image_cost_calculator(
|
||||
model=model,
|
||||
image_response=completion_response,
|
||||
size=resolved_size,
|
||||
n=resolved_n,
|
||||
optional_params=optional_params,
|
||||
)
|
||||
elif custom_llm_provider == litellm.LlmProviders.FAL_AI.value:
|
||||
from litellm.llms.fal_ai.cost_calculator import (
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
import base64
|
||||
from collections.abc import Mapping, Sequence
|
||||
from io import BufferedReader
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
|
||||
from httpx._types import RequestFiles
|
||||
|
|
@ -24,21 +26,12 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
|
|||
Azure AI Foundry FLUX 2 image edit config
|
||||
|
||||
Supports FLUX 2 models (e.g., flux.2-pro) for image editing.
|
||||
Uses the same /providers/blackforestlabs/v1/flux-2-pro endpoint as image generation,
|
||||
Uses the model-specific /providers/blackforestlabs/v1/flux-2-* endpoint as image generation,
|
||||
with the image passed as base64 in JSON body.
|
||||
"""
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
"""
|
||||
FLUX 2 supports a subset of OpenAI image edit params
|
||||
"""
|
||||
return [
|
||||
"prompt",
|
||||
"image",
|
||||
"model",
|
||||
"n",
|
||||
"size",
|
||||
]
|
||||
return AzureFoundryFluxImageGenerationConfig().get_supported_openai_params(model)
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
|
|
@ -50,14 +43,14 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
|
|||
Map OpenAI params to FLUX 2 params.
|
||||
FLUX 2 uses the same param names as OpenAI for supported params.
|
||||
"""
|
||||
mapped_params: Final[dict[str, Any]] = {}
|
||||
supported_params: Final = self.get_supported_openai_params(model)
|
||||
|
||||
for key, value in dict(image_edit_optional_params).items():
|
||||
if key in supported_params and value is not None:
|
||||
mapped_params[key] = value
|
||||
|
||||
return mapped_params
|
||||
return AzureFoundryFluxImageGenerationConfig().map_openai_params(
|
||||
non_default_params=MappingProxyType(
|
||||
{key: value for key, value in image_edit_optional_params.items() if value is not None}
|
||||
),
|
||||
optional_params=MappingProxyType({}),
|
||||
model=model,
|
||||
drop_params=drop_params,
|
||||
)
|
||||
|
||||
def use_multipart_form_data(self) -> bool:
|
||||
"""FLUX 2 uses JSON requests, not multipart/form-data."""
|
||||
|
|
@ -90,7 +83,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
|
|||
self,
|
||||
model: str,
|
||||
prompt: str | None,
|
||||
image: FileTypes | None,
|
||||
image: FileTypes | Sequence[FileTypes] | None,
|
||||
image_edit_optional_request_params: dict,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict,
|
||||
|
|
@ -107,29 +100,29 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
|
|||
if image is None:
|
||||
raise ValueError("FLUX 2 image edit requires an image.")
|
||||
|
||||
image_b64: Final = self._convert_image_to_base64(image)
|
||||
images: Final = tuple(image) if isinstance(image, list) else (image,)
|
||||
if not images:
|
||||
raise ValueError("FLUX 2 image edit requires at least one image.")
|
||||
max_reference_images: Final = 10 if "flex" in model.lower() else 8
|
||||
if len(images) > max_reference_images:
|
||||
raise ValueError(f"{model} supports at most {max_reference_images} reference images.")
|
||||
|
||||
# Build request body with required params
|
||||
reference_images: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"input_image" if index == 1 else f"input_image_{index}": self._convert_image_to_base64(reference_image)
|
||||
for index, reference_image in enumerate(images, start=1)
|
||||
}
|
||||
)
|
||||
request_body: Final[dict[str, Any]] = {
|
||||
"prompt": prompt,
|
||||
"image": image_b64,
|
||||
"model": model,
|
||||
**reference_images,
|
||||
**image_edit_optional_request_params,
|
||||
}
|
||||
|
||||
# Add mapped optional params (already filtered by map_openai_params)
|
||||
request_body.update(image_edit_optional_request_params)
|
||||
|
||||
# Return JSON body and empty files list (FLUX 2 doesn't use multipart)
|
||||
return request_body, []
|
||||
|
||||
def _convert_image_to_base64(self, image: Any) -> str:
|
||||
"""Convert image file to base64 string"""
|
||||
# Handle list of images (take first one)
|
||||
if isinstance(image, list):
|
||||
if len(image) == 0:
|
||||
raise ValueError("Empty image list provided")
|
||||
image = image[0]
|
||||
|
||||
if isinstance(image, BufferedReader):
|
||||
image_bytes = image.read()
|
||||
image.seek(0) # Reset file pointer for potential reuse
|
||||
|
|
@ -151,7 +144,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
|
|||
"""
|
||||
Constructs a complete URL for Azure AI Foundry FLUX 2 image edits.
|
||||
|
||||
Uses the same /providers/blackforestlabs/v1/flux-2-pro endpoint as image generation.
|
||||
Uses the same model-specific BFL provider endpoint as image generation.
|
||||
"""
|
||||
api_base = AzureFoundryModelInfo.get_api_base(api_base)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
import litellm
|
||||
|
|
@ -10,6 +11,9 @@ from litellm.types.utils import ImageResponse
|
|||
def cost_calculator(
|
||||
model: str,
|
||||
image_response: Any,
|
||||
size: str | None = None,
|
||||
n: int | None = None,
|
||||
optional_params: Mapping[str, object] | None = None,
|
||||
) -> float:
|
||||
"""
|
||||
Azure AI image generation cost calculator
|
||||
|
|
@ -28,10 +32,29 @@ def cost_calculator(
|
|||
if token_based_cost is not None:
|
||||
return token_based_cost
|
||||
|
||||
num_images: Final = n if n is not None else len(image_response.data or ())
|
||||
output_cost_per_image: Final[float] = _model_info.get("output_cost_per_image") or 0.0
|
||||
num_images: int = 0
|
||||
if image_response.data:
|
||||
num_images = len(image_response.data)
|
||||
return output_cost_per_image * num_images
|
||||
if output_cost_per_image:
|
||||
return output_cost_per_image * num_images
|
||||
|
||||
model_cost: Final = litellm.model_cost[_model_info["key"]]
|
||||
input_cost_per_pixel: Final[float] = model_cost.get("input_cost_per_pixel") or 0.0
|
||||
if input_cost_per_pixel:
|
||||
from litellm.cost_calculator import default_image_cost_calculator
|
||||
|
||||
width: Final = optional_params.get("width") if optional_params else None
|
||||
height: Final = optional_params.get("height") if optional_params else None
|
||||
pixel_size: Final = (
|
||||
f"{width}x{height}"
|
||||
if type(width) is int and type(height) is int and width > 0 and height > 0
|
||||
else size or image_response.size
|
||||
)
|
||||
return default_image_cost_calculator(
|
||||
model=_model_info["key"],
|
||||
custom_llm_provider=litellm.LlmProviders.AZURE_AI.value,
|
||||
size=pixel_size,
|
||||
n=num_images,
|
||||
)
|
||||
return 0.0
|
||||
|
||||
raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}")
|
||||
|
|
|
|||
|
|
@ -1,18 +1,22 @@
|
|||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from litellm.exceptions import BadRequestError, UnsupportedParamsError
|
||||
from litellm.llms.openai.image_generation import GPTImageGenerationConfig
|
||||
from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
|
||||
|
||||
FLUX2_DROPPED_OPENAI_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = (
|
||||
"background",
|
||||
"moderation",
|
||||
"output_compression",
|
||||
"quality",
|
||||
"user",
|
||||
)
|
||||
|
||||
|
||||
class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
|
||||
"""
|
||||
Azure Foundry flux image generation config
|
||||
|
||||
From manual testing it follows the gpt-image-1 image generation config
|
||||
|
||||
(Azure Foundry does not have any docs on supported params at the time of writing)
|
||||
|
||||
From our test suite - following GPTImageGenerationConfig is working for this model
|
||||
"""
|
||||
"""Azure Foundry BFL API configuration for FLUX image generation."""
|
||||
|
||||
@staticmethod
|
||||
def get_flux2_image_generation_url(
|
||||
|
|
@ -25,11 +29,11 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
|
|||
|
||||
FLUX 2 models on Azure AI use a different URL pattern than standard Azure OpenAI:
|
||||
- Standard: /openai/deployments/{model}/images/generations
|
||||
- FLUX 2: /providers/blackforestlabs/v1/flux-2-pro
|
||||
- FLUX 2: /providers/blackforestlabs/v1/{model-path}
|
||||
|
||||
Args:
|
||||
api_base: Base URL (e.g., https://litellm-ci-cd-prod.services.ai.azure.com)
|
||||
model: Model name (e.g., flux.2-pro)
|
||||
model: Model name (e.g., FLUX.2-flex or FLUX.2-pro)
|
||||
api_version: API version (e.g., preview)
|
||||
|
||||
Returns:
|
||||
|
|
@ -47,9 +51,8 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
|
|||
return api_base
|
||||
return f"{api_base}?api-version={api_version}"
|
||||
|
||||
# Construct the FLUX 2 provider path
|
||||
# Model name flux.2-pro maps to endpoint flux-2-pro
|
||||
return f"{api_base}/providers/blackforestlabs/v1/flux-2-pro?api-version={api_version}"
|
||||
provider_model_path: Final = AzureFoundryFluxImageGenerationConfig.get_flux2_provider_model_path(model)
|
||||
return f"{api_base}/providers/blackforestlabs/v1/{provider_model_path}?api-version={api_version}"
|
||||
|
||||
@staticmethod
|
||||
def is_flux2_model(model: str) -> bool:
|
||||
|
|
@ -64,3 +67,90 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
|
|||
"""
|
||||
model_lower: Final = model.lower().replace(".", "-").replace("_", "-")
|
||||
return "flux-2" in model_lower or "flux2" in model_lower
|
||||
|
||||
@staticmethod
|
||||
def get_flux2_provider_model_path(model: str) -> str:
|
||||
normalized_model: Final = model.lower().replace(".", "-").replace("_", "-")
|
||||
return "flux-2-flex" if "flux-2-flex" in normalized_model else "flux-2-pro"
|
||||
|
||||
def get_supported_openai_params( # mutable-ok: inherited config contract returns a list
|
||||
self, model: str
|
||||
) -> list[OpenAIImageGenerationOptionalParams]:
|
||||
if not self.is_flux2_model(model):
|
||||
return super().get_supported_openai_params(model)
|
||||
return [ # mutable-ok: BaseImageGenerationConfig requires a list
|
||||
"n",
|
||||
"size",
|
||||
"output_format",
|
||||
"seed",
|
||||
"safety_tolerance",
|
||||
"aspect_ratio",
|
||||
"width",
|
||||
"height",
|
||||
"num_images",
|
||||
"guidance",
|
||||
"steps",
|
||||
*FLUX2_DROPPED_OPENAI_PARAMS,
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _map_parameter(name: str, value: object, model: str) -> tuple[tuple[str, object], ...]:
|
||||
if name in FLUX2_DROPPED_OPENAI_PARAMS:
|
||||
return ()
|
||||
if isinstance(value, str):
|
||||
if name in ("n", "num_images", "width", "height", "steps", "seed", "safety_tolerance"):
|
||||
return (("num_images" if name == "n" else name, int(value)),)
|
||||
if name == "guidance":
|
||||
return ((name, float(value)),)
|
||||
if name == "n":
|
||||
return (("num_images", value),)
|
||||
if name != "size":
|
||||
return ((name, value),)
|
||||
if str(value).lower() == "auto":
|
||||
return ()
|
||||
|
||||
try:
|
||||
width, height = (int(dimension) for dimension in str(value).lower().split("x"))
|
||||
except (TypeError, ValueError):
|
||||
raise BadRequestError(
|
||||
message=f"Invalid size format '{value}'. Expected 'WxH', for example '1024x1024'.",
|
||||
model=model,
|
||||
llm_provider="azure_ai",
|
||||
)
|
||||
return (("width", width), ("height", height))
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: Mapping[str, object],
|
||||
optional_params: Mapping[str, object],
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict[str, object]: # mutable-ok: inherited config contract returns a dict
|
||||
if not self.is_flux2_model(model):
|
||||
return super().map_openai_params(
|
||||
non_default_params=dict(non_default_params),
|
||||
optional_params=dict(optional_params),
|
||||
model=model,
|
||||
drop_params=drop_params,
|
||||
)
|
||||
supported_params: Final = self.get_supported_openai_params(model)
|
||||
unsupported_params: Final = tuple(name for name in non_default_params if name not in supported_params)
|
||||
if unsupported_params and not drop_params:
|
||||
raise UnsupportedParamsError(
|
||||
message=(
|
||||
f"Parameters {unsupported_params} are not supported for model {model}. "
|
||||
f"Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters."
|
||||
),
|
||||
model=model,
|
||||
llm_provider="azure_ai",
|
||||
)
|
||||
|
||||
mapped_params: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{
|
||||
mapped_name: mapped_value
|
||||
for name, value in non_default_params.items()
|
||||
if name in supported_params
|
||||
for mapped_name, mapped_value in self._map_parameter(name, value, model)
|
||||
}
|
||||
)
|
||||
return {**optional_params, **mapped_params} # mutable-ok: inherited config contract returns a dict
|
||||
|
|
|
|||
|
|
@ -82,7 +82,14 @@ class GPTImageGenerationConfig(BaseImageGenerationConfig):
|
|||
)
|
||||
|
||||
# set optional params
|
||||
image_response.size = image_response.size or optional_params.get("size", "1024x1024")
|
||||
width: Final = optional_params.get("width")
|
||||
height: Final = optional_params.get("height")
|
||||
requested_size: Final = (
|
||||
f"{width}x{height}"
|
||||
if isinstance(width, int) and isinstance(height, int)
|
||||
else optional_params.get("size", "1024x1024")
|
||||
)
|
||||
image_response.size = image_response.size or requested_size
|
||||
image_response.quality = image_response.quality or optional_params.get("quality", "high")
|
||||
image_response.output_format = image_response.output_format or optional_params.get("output_format", "png")
|
||||
|
||||
|
|
|
|||
|
|
@ -10820,6 +10820,25 @@
|
|||
"/v1/images/generations"
|
||||
]
|
||||
},
|
||||
"azure_ai/FLUX.2-flex": {
|
||||
"input_cost_per_pixel": 5e-08,
|
||||
"litellm_provider": "azure_ai",
|
||||
"max_input_tokens": 32000,
|
||||
"max_tokens": 32000,
|
||||
"mode": "image_generation",
|
||||
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/",
|
||||
"supported_endpoints": [
|
||||
"/v1/images/generations",
|
||||
"/v1/images/edits"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"image"
|
||||
]
|
||||
},
|
||||
"azure_ai/FW-DeepSeek-V3.2": {
|
||||
"deprecation_date": "2027-07-01",
|
||||
"cache_read_input_token_cost": 3.1e-07,
|
||||
|
|
|
|||
|
|
@ -1160,6 +1160,10 @@ OpenAIImageGenerationOptionalParams = Literal[
|
|||
"image_url",
|
||||
"image_prompt_strength",
|
||||
"aspect_ratio",
|
||||
"width",
|
||||
"height",
|
||||
"guidance",
|
||||
"steps",
|
||||
"imageConfig",
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -10820,6 +10820,25 @@
|
|||
"/v1/images/generations"
|
||||
]
|
||||
},
|
||||
"azure_ai/FLUX.2-flex": {
|
||||
"input_cost_per_pixel": 5e-08,
|
||||
"litellm_provider": "azure_ai",
|
||||
"max_input_tokens": 32000,
|
||||
"max_tokens": 32000,
|
||||
"mode": "image_generation",
|
||||
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/",
|
||||
"supported_endpoints": [
|
||||
"/v1/images/generations",
|
||||
"/v1/images/edits"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"image"
|
||||
]
|
||||
},
|
||||
"azure_ai/FW-DeepSeek-V3.2": {
|
||||
"deprecation_date": "2027-07-01",
|
||||
"cache_read_input_token_cost": 3.1e-07,
|
||||
|
|
|
|||
|
|
@ -1,12 +1,20 @@
|
|||
import base64
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.images.utils import ImageEditRequestUtils
|
||||
from litellm.llms.azure_ai.image_edit.flux2_transformation import (
|
||||
AzureFoundryFlux2ImageEditConfig,
|
||||
)
|
||||
from litellm.llms.azure_ai.image_edit.transformation import (
|
||||
AzureFoundryFluxImageEditConfig,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
|
||||
def test_azure_ai_validate_environment():
|
||||
|
|
@ -60,3 +68,122 @@ def test_flux2_validate_environment_with_entra_token(monkeypatch):
|
|||
|
||||
assert headers["Authorization"] == "Bearer entra-token"
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
|
||||
|
||||
def test_flux2_image_edit_maps_openai_and_provider_parameters():
|
||||
config = AzureFoundryFlux2ImageEditConfig()
|
||||
requested_params = ImageEditRequestUtils.get_requested_image_edit_optional_param(
|
||||
{
|
||||
"n": 2,
|
||||
"size": "1536x1024",
|
||||
"guidance": 4.5,
|
||||
"steps": 32,
|
||||
"unrelated": "discarded",
|
||||
},
|
||||
provider_supported_params=config.get_supported_openai_params("FLUX.2-flex"),
|
||||
)
|
||||
mapped_params = config.map_openai_params(
|
||||
image_edit_optional_params=requested_params,
|
||||
model="FLUX.2-flex",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert mapped_params == {
|
||||
"num_images": 2,
|
||||
"width": 1536,
|
||||
"height": 1024,
|
||||
"guidance": 4.5,
|
||||
"steps": 32,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "max_reference_images"),
|
||||
[
|
||||
("FLUX.2-flex", 10),
|
||||
("FLUX.2-pro", 8),
|
||||
],
|
||||
)
|
||||
def test_flux2_image_edit_uses_all_reference_fields(model: str, max_reference_images: int):
|
||||
images = [f"image-{index}".encode() for index in range(1, max_reference_images + 1)]
|
||||
request, files = AzureFoundryFlux2ImageEditConfig().transform_image_edit_request(
|
||||
model=model,
|
||||
prompt="Blend every reference",
|
||||
image=images,
|
||||
image_edit_optional_request_params={"guidance": 4.5, "steps": 20},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert files == []
|
||||
assert request["input_image"] == base64.b64encode(images[0]).decode()
|
||||
assert request[f"input_image_{max_reference_images}"] == base64.b64encode(images[-1]).decode()
|
||||
assert "input_image_1" not in request
|
||||
assert "image" not in request
|
||||
assert len([key for key in request if key.startswith("input_image")]) == max_reference_images
|
||||
assert request["guidance"] == 4.5
|
||||
assert request["steps"] == 20
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "reference_images"),
|
||||
[
|
||||
("FLUX.2-flex", 11),
|
||||
("FLUX.2-pro", 9),
|
||||
],
|
||||
)
|
||||
def test_flux2_image_edit_rejects_too_many_references(model: str, reference_images: int):
|
||||
with pytest.raises(ValueError, match=f"at most {reference_images - 1} reference images"):
|
||||
AzureFoundryFlux2ImageEditConfig().transform_image_edit_request(
|
||||
model=model,
|
||||
prompt="Blend every reference",
|
||||
image=[b"image"] * reference_images,
|
||||
image_edit_optional_request_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dimensions", ({"size": "2048x1024"}, {"width": 2048, "height": 1024}, {"width": "2048", "height": "1024"}))
|
||||
@pytest.mark.usefixtures("local_model_cost_map")
|
||||
def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[str, int | str]):
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
body: Final = json.loads(request.content)
|
||||
assert body == {
|
||||
"model": "FLUX.2-flex",
|
||||
"prompt": "Add a hat",
|
||||
"input_image": base64.b64encode(b"image").decode(),
|
||||
"num_images": 2,
|
||||
"width": 2048,
|
||||
"height": 1024,
|
||||
"guidance": 4.5,
|
||||
"steps": 32,
|
||||
}
|
||||
return httpx.Response(200, json={"data": [{"b64_json": "aW1n"}, {"b64_json": "aW1n"}]})
|
||||
|
||||
client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond)))
|
||||
response: Final = litellm.image_edit(
|
||||
model="azure_ai/FLUX.2-flex",
|
||||
image=b"image",
|
||||
prompt="Add a hat",
|
||||
api_key="test-key",
|
||||
api_base="https://example.services.ai.azure.com",
|
||||
client=client,
|
||||
n=2,
|
||||
guidance="4.5",
|
||||
steps="32",
|
||||
**dimensions,
|
||||
)
|
||||
|
||||
assert response._hidden_params["response_cost"] == pytest.approx(5e-08 * 2048 * 1024 * 2)
|
||||
|
||||
|
||||
def test_flux2_image_edit_accepts_and_drops_openai_only_parameters():
|
||||
optional_params: Final = ImageEditRequestUtils.get_optional_params_image_edit(
|
||||
model="FLUX.2-pro",
|
||||
image_edit_provider_config=AzureFoundryFlux2ImageEditConfig(),
|
||||
image_edit_optional_params={"n": 1, "size": "auto", "quality": "high", "user": "end-user-1"},
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert optional_params == {"num_images": 1}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,223 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils
|
||||
from litellm.llms.azure.azure import AzureChatCompletion
|
||||
from litellm.llms.azure.image_generation import get_azure_image_generation_config
|
||||
from litellm.llms.azure.image_generation.http_utils import azure_deployment_image_generation_json_body
|
||||
from litellm.llms.azure_ai.image_generation.flux_transformation import (
|
||||
AzureFoundryFluxImageGenerationConfig,
|
||||
)
|
||||
from litellm.types.utils import ImageObject, ImageResponse
|
||||
from litellm.utils import _invalidate_model_cost_lowercase_map, get_optional_params_image_gen
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def use_local_model_cost_map(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map())
|
||||
litellm.get_model_info.cache_clear()
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
yield
|
||||
litellm.get_model_info.cache_clear()
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "provider_path"),
|
||||
[
|
||||
("FLUX.2-flex", "flux-2-flex"),
|
||||
("FLUX.2-pro", "flux-2-pro"),
|
||||
],
|
||||
)
|
||||
def test_flux2_uses_model_specific_provider_url(model: str, provider_path: str):
|
||||
url = AzureChatCompletion().create_azure_base_url(
|
||||
azure_client_params={
|
||||
"azure_endpoint": "https://example.services.ai.azure.com/",
|
||||
"api_version": "preview",
|
||||
},
|
||||
model=model,
|
||||
)
|
||||
|
||||
assert (
|
||||
url == f"https://example.services.ai.azure.com/providers/blackforestlabs/v1/{provider_path}?api-version=preview"
|
||||
)
|
||||
|
||||
|
||||
def test_flux2_flex_maps_openai_and_provider_parameters():
|
||||
config = AzureFoundryFluxImageGenerationConfig()
|
||||
mapped_params = config.map_openai_params(
|
||||
non_default_params={
|
||||
"n": 2,
|
||||
"size": "1536x1024",
|
||||
"guidance": 4.5,
|
||||
"steps": 32,
|
||||
"output_format": "jpeg",
|
||||
},
|
||||
optional_params={},
|
||||
model="FLUX.2-flex",
|
||||
drop_params=False,
|
||||
)
|
||||
url = config.get_flux2_image_generation_url(
|
||||
api_base="https://example.services.ai.azure.com",
|
||||
model="FLUX.2-flex",
|
||||
api_version="preview",
|
||||
)
|
||||
request = azure_deployment_image_generation_json_body(
|
||||
api_base=url,
|
||||
data={"model": "FLUX.2-flex", "prompt": "A red fox", **mapped_params},
|
||||
deployment_name="FLUX.2-flex",
|
||||
)
|
||||
|
||||
assert request == {
|
||||
"model": "FLUX.2-flex",
|
||||
"prompt": "A red fox",
|
||||
"num_images": 2,
|
||||
"width": 1536,
|
||||
"height": 1024,
|
||||
"guidance": 4.5,
|
||||
"steps": 32,
|
||||
"output_format": "jpeg",
|
||||
}
|
||||
|
||||
|
||||
def test_flux2_flex_rejects_invalid_size_as_bad_request():
|
||||
with pytest.raises(litellm.BadRequestError, match="Expected 'WxH'") as raised:
|
||||
get_optional_params_image_gen(
|
||||
model="FLUX.2-flex",
|
||||
custom_llm_provider="azure_ai",
|
||||
provider_config=AzureFoundryFluxImageGenerationConfig(),
|
||||
size="large",
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ("FLUX.2-pro", "FLUX.2-flex"))
|
||||
def test_flux2_accepts_and_drops_openai_only_image_parameters(model: str):
|
||||
optional_params: Final = get_optional_params_image_gen(
|
||||
model=model,
|
||||
custom_llm_provider="azure_ai",
|
||||
provider_config=AzureFoundryFluxImageGenerationConfig(),
|
||||
n=1,
|
||||
size="auto",
|
||||
quality="high",
|
||||
user="end-user-1",
|
||||
background="transparent",
|
||||
moderation="low",
|
||||
output_compression=50,
|
||||
)
|
||||
|
||||
assert optional_params == {"num_images": 1}
|
||||
|
||||
|
||||
def test_flux2_flex_model_info():
|
||||
model_info = litellm.get_model_info(
|
||||
model="FLUX.2-flex",
|
||||
custom_llm_provider="azure_ai",
|
||||
)
|
||||
catalog_info = litellm.model_cost["azure_ai/FLUX.2-flex"]
|
||||
|
||||
assert model_info["mode"] == "image_generation"
|
||||
assert model_info["max_input_tokens"] == 32000
|
||||
assert model_info["max_tokens"] == 32000
|
||||
assert model_info["supported_endpoints"] == ["/v1/images/generations", "/v1/images/edits"]
|
||||
assert catalog_info["input_cost_per_pixel"] == 5e-08
|
||||
assert catalog_info["supported_modalities"] == ["text", "image"]
|
||||
assert catalog_info["supported_output_modalities"] == ["image"]
|
||||
|
||||
|
||||
def test_flux2_flex_cost_uses_generated_megapixels():
|
||||
response = ImageResponse(
|
||||
data=[
|
||||
ImageObject(url="https://example.com/one.png"),
|
||||
ImageObject(url="https://example.com/two.png"),
|
||||
]
|
||||
)
|
||||
|
||||
cost = CostCalculatorUtils.route_image_generation_cost_calculator(
|
||||
model="FLUX.2-flex",
|
||||
completion_response=response,
|
||||
custom_llm_provider="azure_ai",
|
||||
size="2048x1024",
|
||||
call_type="image_generation",
|
||||
)
|
||||
|
||||
assert cost == pytest.approx(5e-08 * 2048 * 1024 * 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ("FLUX-1.1-pro", "FLUX.1-Kontext-pro"))
|
||||
def test_flux1_preserves_existing_openai_parameters(model: str):
|
||||
params: Final = {"n": 2, "size": "1536x1024", "quality": "high", "user": "test-user"}
|
||||
|
||||
mapped: Final = AzureFoundryFluxImageGenerationConfig().map_openai_params(
|
||||
non_default_params=params,
|
||||
optional_params={},
|
||||
model=model,
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert mapped == params
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dimensions", ({"size": "2048x1024"}, {"width": 2048, "height": 1024}))
|
||||
def test_flux2_cost_uses_mapped_dimensions_after_response_transformation(dimensions: Mapping[str, int | str]):
|
||||
params: Final = AzureFoundryFluxImageGenerationConfig().map_openai_params(
|
||||
non_default_params={"n": 2, **dimensions},
|
||||
optional_params={},
|
||||
model="FLUX.2-flex",
|
||||
drop_params=False,
|
||||
)
|
||||
response: Final = get_azure_image_generation_config("FLUX.2-flex").transform_image_generation_response(
|
||||
model="FLUX.2-flex",
|
||||
raw_response=httpx.Response(200, json={"data": [{"b64_json": "aW1n"}, {"b64_json": "aW1n"}]}),
|
||||
model_response=ImageResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
request_data={"prompt": "A red fox", **params},
|
||||
optional_params=params,
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
assert litellm.completion_cost(
|
||||
model="azure_ai/FLUX.2-flex",
|
||||
completion_response=response,
|
||||
optional_params=params,
|
||||
call_type="image_generation",
|
||||
) == pytest.approx(5e-08 * 2048 * 1024 * 2)
|
||||
|
||||
|
||||
def test_flux2_flex_cost_accepts_lowercase_model_spelling():
|
||||
response: Final = ImageResponse(data=[ImageObject(b64_json="aW1n"), ImageObject(b64_json="aW1n")])
|
||||
|
||||
cost: Final = litellm.completion_cost(
|
||||
model="azure_ai/flux.2-flex",
|
||||
completion_response=response,
|
||||
optional_params={"width": 1536, "height": 1024, "num_images": 2},
|
||||
call_type="image_generation",
|
||||
)
|
||||
|
||||
assert cost == pytest.approx(5e-08 * 1536 * 1024 * 2)
|
||||
|
||||
|
||||
def test_flux2_response_preserves_mapped_dimensions():
|
||||
config = AzureFoundryFluxImageGenerationConfig()
|
||||
params = config.map_openai_params(
|
||||
non_default_params={"size": "2048x1024"}, optional_params={}, model="FLUX.2-flex", drop_params=False
|
||||
)
|
||||
response = config.transform_image_generation_response(
|
||||
model="FLUX.2-flex",
|
||||
raw_response=httpx.Response(200, json={"data": [{"b64_json": "aW1n"}]}),
|
||||
model_response=ImageResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
request_data={"prompt": "A landscape"},
|
||||
optional_params=params,
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
assert response.size == "2048x1024"
|
||||
Loading…
Add table
Reference in a new issue