From 4595b4f62f209d3d71734e8f1f2692f9a726e9a1 Mon Sep 17 00:00:00 2001
From: Emerson Gomes <emerson.gomes@thalesgroup.com>
Date: Tue, 15 Sep 2026 13:01:12 -0500
Subject: [PATCH] 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"