From 911f66aff69fabb6666bde3f54db70960cb04b56 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 12:12:57 -0500 Subject: [PATCH 1/6] feat(azure_ai): support FLUX.2 flex images --- litellm/images/main.py | 9 +- litellm/images/utils.py | 7 +- .../litellm_core_utils/llm_cost_calc/utils.py | 3 + .../image_edit/flux2_transformation.py | 64 ++++--- .../image_generation/cost_calculator.py | 34 +++- .../image_generation/flux_transformation.py | 91 +++++++-- ...odel_prices_and_context_window_backup.json | 19 ++ litellm/proxy/_lazy_openapi_snapshot.json | 2 +- litellm/types/llms/openai.py | 4 + model_prices_and_context_window.json | 19 ++ ...test_azure_ai_image_edit_transformation.py | 116 ++++++++++++ .../test_azure_ai_flux2_image_generation.py | 172 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 13 files changed, 491 insertions(+), 53 deletions(-) create mode 100644 tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py diff --git a/litellm/images/main.py b/litellm/images/main.py index 6a94e7c8df2..81547a153c3 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -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 diff --git a/litellm/images/utils.py b/litellm/images/utils.py index 49b70870de6..24454954714 100644 --- a/litellm/images/utils.py +++ b/litellm/images/utils.py @@ -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) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index baa9aab1087..0a0e92ff3a3 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -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 ( diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index a09a80985b7..aa8905e5601 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -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) diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index 106c7e42b83..086293f26c0 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -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)}") diff --git a/litellm/llms/azure_ai/image_generation/flux_transformation.py b/litellm/llms/azure_ai/image_generation/flux_transformation.py index 65b5a35af52..b10e8a6f35e 100644 --- a/litellm/llms/azure_ai/image_generation/flux_transformation.py +++ b/litellm/llms/azure_ai/image_generation/flux_transformation.py @@ -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 diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9f91cf82f41..bcb2dfc53e8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index fb2f014e3d8..81889e4da0c 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -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": { diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index e3eac9b9205..710b34116e5 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1160,6 +1160,10 @@ OpenAIImageGenerationOptionalParams = Literal[ "image_url", "image_prompt_strength", "aspect_ratio", + "width", + "height", + "guidance", + "steps", "imageConfig", ] diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9f91cf82f41..bcb2dfc53e8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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, diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index b6cb7ea9b54..94b23c5f1c2 100644 --- a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -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) diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py new file mode 100644 index 00000000000..07026cf4309 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -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) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b26f5e25b6f..6b7337ea8d1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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; From ff1e2a02ba9b7bafee87896edcbf935749bf59c3 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 12:33:58 -0500 Subject: [PATCH 2/6] fix(proxy): preserve CI-compatible OpenAPI snapshot formatting --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 81889e4da0c..fb2f014e3d8 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -18974,7 +18974,7 @@ } } }, - "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" + "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 " }, "500": { "content": { From 3821e5ace617f43122e757ec0d922c6f4a8c9b6a Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 12:43:24 -0500 Subject: [PATCH 3/6] fix(azure_ai): specify FLUX parameter mapping return type --- litellm/llms/azure_ai/image_generation/flux_transformation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/azure_ai/image_generation/flux_transformation.py b/litellm/llms/azure_ai/image_generation/flux_transformation.py index b10e8a6f35e..b9d5e11ff2a 100644 --- a/litellm/llms/azure_ai/image_generation/flux_transformation.py +++ b/litellm/llms/azure_ai/image_generation/flux_transformation.py @@ -102,7 +102,7 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): optional_params: Mapping[str, object], model: str, drop_params: bool, - ) -> dict: # mutable-ok: inherited config contract returns a dict + ) -> 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), From 4595b4f62f209d3d71734e8f1f2692f9a726e9a1 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 13:01:12 -0500 Subject: [PATCH 4/6] fix(azure-ai): coerce FLUX controls and preserve response dimensions --- .../image_generation/flux_transformation.py | 5 +++++ .../image_generation/gpt_transformation.py | 9 ++++++++- .../test_azure_ai_image_edit_transformation.py | 6 +++--- .../test_azure_ai_flux2_image_generation.py | 18 ++++++++++++++++++ 4 files changed, 34 insertions(+), 4 deletions(-) diff --git a/litellm/llms/azure_ai/image_generation/flux_transformation.py b/litellm/llms/azure_ai/image_generation/flux_transformation.py index b9d5e11ff2a..997205b1fc3 100644 --- a/litellm/llms/azure_ai/image_generation/flux_transformation.py +++ b/litellm/llms/azure_ai/image_generation/flux_transformation.py @@ -85,6 +85,11 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): @staticmethod def _map_parameter(name: str, value: object) -> tuple[tuple[str, object], ...]: + 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": diff --git a/litellm/llms/openai/image_generation/gpt_transformation.py b/litellm/llms/openai/image_generation/gpt_transformation.py index 090b2eba387..8dc4d8953ea 100644 --- a/litellm/llms/openai/image_generation/gpt_transformation.py +++ b/litellm/llms/openai/image_generation/gpt_transformation.py @@ -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") diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index 94b23c5f1c2..d74afa88a6b 100644 --- a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -144,7 +144,7 @@ def test_flux2_image_edit_rejects_too_many_references(model: str, reference_imag ) -@pytest.mark.parametrize("dimensions", ({"size": "2048x1024"}, {"width": 2048, "height": 1024})) +@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: @@ -170,8 +170,8 @@ def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[ api_base="https://example.services.ai.azure.com", client=client, n=2, - guidance=4.5, - steps=32, + guidance="4.5", + steps="32", **dimensions, ) diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index 07026cf4309..c1d46b7e919 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -170,3 +170,21 @@ def test_flux2_cost_uses_mapped_dimensions_after_response_transformation(dimensi optional_params=params, call_type="image_generation", ) == pytest.approx(5e-08 * 2048 * 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" From cf08440557bd0a377f1ce082497d037294c6aca0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 11:42:52 -0700 Subject: [PATCH 5/6] fix(azure-ai): price FLUX.2 Flex by its resolved cost map key --- .../azure_ai/image_generation/cost_calculator.py | 5 +---- .../test_azure_ai_flux2_image_generation.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index 086293f26c0..35d0f4fb6c3 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -42,9 +42,6 @@ def cost_calculator( 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 = ( @@ -53,7 +50,7 @@ def cost_calculator( else size or image_response.size ) return default_image_cost_calculator( - model=cost_model, + model=_model_info["key"], custom_llm_provider=litellm.LlmProviders.AZURE_AI.value, size=pixel_size, n=num_images, diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index c1d46b7e919..388cca9f054 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -172,6 +172,19 @@ def test_flux2_cost_uses_mapped_dimensions_after_response_transformation(dimensi ) == 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( From c0b0ba20b23a83a5e382ee2b7de7739cf5dcaa10 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 12:41:40 -0700 Subject: [PATCH 6/6] fix(azure-ai): keep FLUX.2 tolerant of OpenAI-only image params --- .../image_edit/flux2_transformation.py | 17 +--------- .../image_generation/flux_transformation.py | 34 +++++++++++++++---- ...test_azure_ai_image_edit_transformation.py | 11 ++++++ .../test_azure_ai_flux2_image_generation.py | 34 +++++++++++++++---- 4 files changed, 67 insertions(+), 29 deletions(-) diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index aa8905e5601..f91a87ba0f4 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -31,22 +31,7 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): """ def get_supported_openai_params(self, model: str) -> list: - """ - FLUX 2 supports a subset of OpenAI image edit params - """ - return [ - "n", - "size", - "width", - "height", - "num_images", - "seed", - "safety_tolerance", - "output_format", - "aspect_ratio", - "guidance", - "steps", - ] + return AzureFoundryFluxImageGenerationConfig().get_supported_openai_params(model) def map_openai_params( self, diff --git a/litellm/llms/azure_ai/image_generation/flux_transformation.py b/litellm/llms/azure_ai/image_generation/flux_transformation.py index 997205b1fc3..ac9ec24420b 100644 --- a/litellm/llms/azure_ai/image_generation/flux_transformation.py +++ b/litellm/llms/azure_ai/image_generation/flux_transformation.py @@ -2,9 +2,18 @@ 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 BFL API configuration for FLUX image generation.""" @@ -81,10 +90,13 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): "num_images", "guidance", "steps", + *FLUX2_DROPPED_OPENAI_PARAMS, ] @staticmethod - def _map_parameter(name: str, value: object) -> tuple[tuple[str, object], ...]: + 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)),) @@ -94,11 +106,17 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): 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 ValueError(f"Invalid size format '{value}'. Expected 'WxH', for example '1024x1024'.") + 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( @@ -118,9 +136,13 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): 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." + 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( @@ -128,7 +150,7 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): 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) + 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 diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index d74afa88a6b..39001c1795b 100644 --- a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -176,3 +176,14 @@ def test_flux2_image_edit_preserves_controls_and_pixel_cost(dimensions: Mapping[ ) 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} diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index 388cca9f054..512e98b4151 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -15,7 +15,7 @@ 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 +from litellm.utils import _invalidate_model_cost_lowercase_map, get_optional_params_image_gen @pytest.fixture(autouse=True) @@ -86,15 +86,35 @@ def test_flux2_flex_maps_openai_and_provider_parameters(): } -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={}, +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", - drop_params=False, + 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(