mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
feat(azure_ai): support FLUX.2 flex images
This commit is contained in:
Emerson Gomes
parent
3ad9a7f336
commit
911f66aff6
13 changed files with 491 additions and 53 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)
|
||||
|
||||
|
|
|
|||
|
|
@ -1853,6 +1853,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,7 +26,7 @@ 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.
|
||||
"""
|
||||
|
||||
|
|
@ -33,11 +35,17 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
|
|||
FLUX 2 supports a subset of OpenAI image edit params
|
||||
"""
|
||||
return [
|
||||
"prompt",
|
||||
"image",
|
||||
"model",
|
||||
"n",
|
||||
"size",
|
||||
"width",
|
||||
"height",
|
||||
"num_images",
|
||||
"seed",
|
||||
"safety_tolerance",
|
||||
"output_format",
|
||||
"aspect_ratio",
|
||||
"guidance",
|
||||
"steps",
|
||||
]
|
||||
|
||||
def map_openai_params(
|
||||
|
|
@ -50,14 +58,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 +98,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 +115,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 +159,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,32 @@ 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
|
||||
|
||||
cost_model: Final = (
|
||||
model if model.startswith(f"{litellm.LlmProviders.AZURE_AI.value}/") else f"azure_ai/{model}"
|
||||
)
|
||||
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=cost_model,
|
||||
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,13 @@
|
|||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from litellm.llms.openai.image_generation import GPTImageGenerationConfig
|
||||
from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
|
||||
|
||||
|
||||
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 +20,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 +42,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 +58,72 @@ 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",
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _map_parameter(name: str, value: object) -> tuple[tuple[str, object], ...]:
|
||||
if name == "n":
|
||||
return (("num_images", value),)
|
||||
if name != "size":
|
||||
return ((name, value),)
|
||||
|
||||
try:
|
||||
width, height = (int(dimension) for dimension in str(value).lower().split("x"))
|
||||
except (TypeError, ValueError):
|
||||
raise ValueError(f"Invalid size format '{value}'. Expected 'WxH', for example '1024x1024'.")
|
||||
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: # 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 ValueError(
|
||||
f"Parameters {unsupported_params} are not supported for model {model}. "
|
||||
f"Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters."
|
||||
)
|
||||
|
||||
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)
|
||||
}
|
||||
)
|
||||
return {**optional_params, **mapped_params} # mutable-ok: inherited config contract returns a dict
|
||||
|
|
|
|||
|
|
@ -9900,6 +9900,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,
|
||||
|
|
|
|||
|
|
@ -18974,7 +18974,7 @@
|
|||
}
|
||||
}
|
||||
},
|
||||
"description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n "
|
||||
"description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n"
|
||||
},
|
||||
"500": {
|
||||
"content": {
|
||||
|
|
|
|||
|
|
@ -1160,6 +1160,10 @@ OpenAIImageGenerationOptionalParams = Literal[
|
|||
"image_url",
|
||||
"image_prompt_strength",
|
||||
"aspect_ratio",
|
||||
"width",
|
||||
"height",
|
||||
"guidance",
|
||||
"steps",
|
||||
"imageConfig",
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -9900,6 +9900,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,111 @@ 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}))
|
||||
@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)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,172 @@
|
|||
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
|
||||
|
||||
|
||||
@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():
|
||||
with pytest.raises(ValueError, match="Expected 'WxH'"):
|
||||
AzureFoundryFluxImageGenerationConfig().map_openai_params(
|
||||
non_default_params={"size": "large"},
|
||||
optional_params={},
|
||||
model="FLUX.2-flex",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
|
||||
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)
|
||||
4
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
4
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -35343,7 +35343,7 @@ export interface components {
|
|||
default_model?: string | null;
|
||||
/**
|
||||
* Deployment Affinity
|
||||
* @description When True and a session_id is resolvable on the request, pin the deployment chosen inside each routed model group and reuse it whenever the session returns to that group, without pinning which group the session routes to. Independent of session_affinity, which pins the model group instead (and always carries this deployment pin with it): with session_affinity off, every turn is still classified on its own merits while a session that escalates to a stronger tier and comes back still lands on the deployment it used before, which is what keeps a provider prompt cache warm. Pins are held per model group, so switching tiers does not disturb the pin left behind in the previous group. On by default because re-shuffling a conversation across deployments of the same model discards that cache for no benefit; set False to keep every turn load-balanced across the group, which is what a deployment set with tight per-deployment rate limits wants. Inert when no session_id is resolvable, since there is nothing to key a pin on, and suppressed when plugins are configured, for the same reason session_affinity is.
|
||||
* @description When True and a client session_id is resolvable, reuse the session's chosen model for each classified tier and its deployment within each model group. With session_affinity off, every turn is still classified: moving to another tier leaves the previous tier's model pin intact for a later return. Pins yield to current candidate, context, modality, and availability constraints. Adaptive selection chooses the initial model from its eligible pool, then reuses that choice per tier. This reduces avoidable provider prompt-cache misses; it does not guarantee cache hits. Set False to select models and load-balance deployments on every turn, unless session_affinity or user_turn classification requires a pin. Inert without a client session_id and suppressed when plugins are configured.
|
||||
* @default true
|
||||
*/
|
||||
deployment_affinity: boolean;
|
||||
|
|
@ -35487,7 +35487,7 @@ export interface components {
|
|||
session_affinity: boolean;
|
||||
/**
|
||||
* Session Affinity Ttl Seconds
|
||||
* @description TTL for the session affinity pin; refreshed on every cache hit. Bounds both the session_affinity model pin and the deployment_affinity deployment pin, so it measures idle time for the session's routing decisions rather than total session length
|
||||
* @description TTL for the session affinity pin; refreshed on every cache hit. Bounds both the session_affinity model pin and the deployment_affinity per-tier model and deployment pins, so it measures idle time for the session's routing decisions rather than total session length
|
||||
* @default 3600
|
||||
*/
|
||||
session_affinity_ttl_seconds: number;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue