diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index 3aac08ddcaf1..f42b99644c7d 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -26,15 +26,14 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): """ 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", + "width", + "height", ] def map_openai_params( @@ -43,15 +42,17 @@ def map_openai_params( model: str, drop_params: bool, ) -> dict: - """ - 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: + if key not in supported_params or value is None: + continue + if key == "size" and isinstance(value, str) and "x" in value: + w, h = value.lower().split("x", 1) + mapped_params["width"] = int(w) + mapped_params["height"] = int(h) + elif key != "size": mapped_params[key] = value return mapped_params 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 da1041f3d607..36c81c33b99d 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 @@ -5,6 +5,9 @@ 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path +from litellm.llms.azure_ai.image_edit.flux2_transformation import ( + AzureFoundryFlux2ImageEditConfig, +) from litellm.llms.azure_ai.image_edit.transformation import ( AzureFoundryFluxImageEditConfig, ) @@ -32,3 +35,24 @@ def test_azure_ai_url_generation(): ) expected_url = f"{api_base}/openai/deployments/FLUX.1-Kontext-pro/images/edits?api-version=2025-04-01-preview" assert complete_url == expected_url + + +def test_flux2_size_param_translates_to_width_height(): + config = AzureFoundryFlux2ImageEditConfig() + result = config.map_openai_params( + image_edit_optional_params={"size": "896x1184"}, + model="flux.2-pro", + drop_params=False, + ) + assert result == {"width": 896, "height": 1184} + assert "size" not in result + + +def test_flux2_explicit_width_height_pass_through(): + config = AzureFoundryFlux2ImageEditConfig() + result = config.map_openai_params( + image_edit_optional_params={"width": 512, "height": 512}, + model="flux.2-pro", + drop_params=False, + ) + assert result == {"width": 512, "height": 512}