feat(azure_ai): support FLUX.2 flex images

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

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 @@ def get_optional_params_image_edit(
     @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 @@ def get_requested_image_edit_optional_param(
         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)
 

litellm/litellm_core_utils/llm_cost_calc/utils.py

@@ -1853,6 +1853,9 @@ def route_image_generation_cost_calculator(
             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 (

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 @@ def map_openai_params(
         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 @@ def transform_image_edit_request(
         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 @@ def transform_image_edit_request(
         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 @@ def get_complete_url(
         """
         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)
 

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 @@
 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)}")

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
-    """
+class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig):
+    """Azure Foundry BFL API configuration for FLUX image generation."""
 
     @staticmethod
     def get_flux2_image_generation_url(
@@ -25,11 +29,11 @@ def get_flux2_image_generation_url(
 
         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 @@ def get_flux2_image_generation_url(
                 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 @@ def is_flux2_model(model: str) -> bool:
         """
         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

litellm/llms/openai/image_generation/gpt_transformation.py

@@ -82,7 +82,14 @@ def transform_image_generation_response(
         )
 
         # 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")