fix(azure_ai): translate size param to width/height for flux.2-pro image edits

## TLDR

Signed-off-by: Ishaan <ishaangupta0408@gmail.com>

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 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
         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

tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py

@@ -5,6 +5,9 @@ sys.path.insert(
     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}