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 6a617d962ea..f8eb15dca88 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -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 ( diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index a09a80985b7..f91a87ba0f4 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,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) diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index 106c7e42b83..35d0f4fb6c3 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,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)}") diff --git a/litellm/llms/azure_ai/image_generation/flux_transformation.py b/litellm/llms/azure_ai/image_generation/flux_transformation.py index 65b5a35af52..ac9ec24420b 100644 --- a/litellm/llms/azure_ai/image_generation/flux_transformation.py +++ b/litellm/llms/azure_ai/image_generation/flux_transformation.py @@ -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 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/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index f8a537fb922..377e30a23b7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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, 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 f8a537fb922..377e30a23b7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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, 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..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 @@ -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} 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..512e98b4151 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -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"