From 38a3de764154a600441a37c746d78c7ca89d57bf Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 8 Sep 2026 18:46:18 -0700
Subject: [PATCH] fix(azure_ai): use the pixel cap the live MAI endpoint
 enforces

---
 .../image_generation/mai_transformation.py    |  2 +-
 .../test_mai_image_generation.py              | 23 ++++++++++++++++++-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py
index 2d9d3246b4c..7f7a1044689 100644
--- a/litellm/llms/azure_ai/image_generation/mai_transformation.py
+++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py
@@ -24,7 +24,7 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
 
     MAX_IMAGES_PER_REQUEST: Final = 1
     MIN_DIMENSION_PX: Final = 768
-    MAX_TOTAL_PX: Final = 1024 * 1024
+    MAX_TOTAL_PX: Final = 1_056_768
 
     @staticmethod
     def get_mai_image_generation_url(
diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
index 29aa5967c9a..a0bc3fb00b7 100644
--- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
+++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
@@ -213,7 +213,7 @@ class TestAzureMAIImageGeneration:
     @pytest.mark.parametrize("size", ["1792x1024", "1024x1792"])
     def test_map_openai_params_size_over_total_pixel_budget_raises(self, size):
         config = AzureFoundryMAIImageGenerationConfig()
-        with pytest.raises(UnsupportedParamsError, match="at most 1048576 total pixels"):
+        with pytest.raises(UnsupportedParamsError, match="at most 1056768 total pixels"):
             config.map_openai_params(
                 non_default_params={"size": size},
                 optional_params={},
@@ -221,6 +221,27 @@ class TestAzureMAIImageGeneration:
                 drop_params=True,
             )
 
+    @pytest.mark.parametrize("size", ["1032x1024", "1376x768"])
+    def test_map_openai_params_size_at_live_pixel_cap_passes_through(self, size):
+        config = AzureFoundryMAIImageGenerationConfig()
+        optional_params = config.map_openai_params(
+            non_default_params={"size": size},
+            optional_params={},
+            model="MAI-Image-2.5",
+            drop_params=False,
+        )
+        assert optional_params["width"] * optional_params["height"] == 1_056_768
+
+    def test_map_openai_params_size_one_pixel_over_live_cap_raises(self):
+        config = AzureFoundryMAIImageGenerationConfig()
+        with pytest.raises(UnsupportedParamsError, match="at most 1056768 total pixels"):
+            config.map_openai_params(
+                non_default_params={"size": "1033x1024"},
+                optional_params={},
+                model="MAI-Image-2.5",
+                drop_params=False,
+            )
+
     def test_map_openai_params_explicit_width_height_not_range_checked(self):
         config = AzureFoundryMAIImageGenerationConfig()
         optional_params = config.map_openai_params(