diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 64f81956ad7..3421cc0a4a1 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -21,6 +21,19 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): DEFAULT_WIDTH = 1024 DEFAULT_HEIGHT = 1024 + # The MAI endpoint produces exactly one image per request. Its documented + # body is model/prompt/width/height (plus `image` for edits) — there is no + # count field, and `n` (or the native `sampleCount`) is accepted and + # ignored, so a request for more silently comes back with one. + MAX_IMAGES_PER_REQUEST: Final = 1 + + # Provider-side bounds on the generated image. Both are enforced by the + # MAI endpoint, which 400s with "'width' must be at least 768 pixels." + # Only `size` is checked against them: `width`/`height` pass through + # unmapped, which keeps a future model with different bounds reachable. + MIN_DIMENSION_PX: Final = 768 + MAX_TOTAL_PX: Final = 1024 * 1024 + @staticmethod def get_mai_image_generation_url( api_base: str | None, @@ -146,6 +159,15 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if k in supported_params: if k == "size" and v: self._map_size_param(v, optional_params) + elif k == "n" and v is not None and v > self.MAX_IMAGES_PER_REQUEST: + if not drop_params: + raise ValueError( + f"n={v} is not supported for model {model}. The Azure AI MAI image " + f"endpoint returns exactly {self.MAX_IMAGES_PER_REQUEST} image per " + "request and ignores any count, so a larger value would silently " + "return fewer images than requested. Send one request per image, or " + "set drop_params=True to drop n." + ) else: optional_params[k] = v elif k in ("width", "height"): @@ -176,13 +198,9 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if size in size_mapping: width, height = size_mapping[size] - optional_params["width"] = width - optional_params["height"] = height elif "x" in size: try: width, height = map(int, size.lower().split("x")) - optional_params["width"] = width - optional_params["height"] = height except ValueError: raise ValueError(f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024').") else: @@ -191,6 +209,29 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string." ) + self._validate_dimensions(size=size, width=width, height=height) + optional_params["width"] = width + optional_params["height"] = height + + def _validate_dimensions(self, size: str, width: int, height: int) -> None: + """Reject a `size` the MAI endpoint would 400 on. + + Several OpenAI-standard sizes are outside MAI's bounds: 512x512 and + 256x256 fall under the per-side minimum, and 1792x1024 / 1024x1792 + exceed the total pixel budget. Checking here turns an opaque provider + 400 into an error that names the constraint. + """ + if width < self.MIN_DIMENSION_PX or height < self.MIN_DIMENSION_PX: + raise ValueError( + f"Unsupported size value: '{size}'. Azure AI MAI image models require width and " + f"height of at least {self.MIN_DIMENSION_PX} pixels." + ) + if width * height > self.MAX_TOTAL_PX: + raise ValueError( + f"Unsupported size value: '{size}'. Azure AI MAI image models accept at most " + f"{self.MAX_TOTAL_PX} total pixels ({width}x{height} is {width * height})." + ) + def transform_image_generation_response( self, model: str, 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 2a44e77ce09..6cbc79fff2b 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 @@ -1,10 +1,8 @@ -import os from unittest.mock import MagicMock import httpx import pytest - import litellm from litellm.llms.azure.azure import AzureChatCompletion from litellm.llms.azure.image_generation import get_azure_image_generation_config @@ -30,9 +28,7 @@ from litellm.utils import get_optional_params_image_gen class TestAzureMAIImageGeneration: def test_is_mai_model(self): assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5") - assert AzureFoundryMAIImageGenerationConfig.is_mai_model( - "azure_ai/MAI-Image-2.5" - ) + assert AzureFoundryMAIImageGenerationConfig.is_mai_model("azure_ai/MAI-Image-2.5") assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5-Flash") assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2e") assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro") @@ -62,16 +58,10 @@ class TestAzureMAIImageGeneration: api_base="https://my-resource.services.ai.azure.com", api_version="preview", ) - assert ( - url - == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" - ) + assert url == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" def test_get_mai_image_generation_url_preserves_full_path(self): - api = ( - "https://my-resource.services.ai.azure.com/mai/v1/images/generations" - "?api-version=preview" - ) + api = "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( api_base=api, api_version="preview", @@ -83,10 +73,7 @@ class TestAzureMAIImageGeneration: api_base="https://my-resource.services.ai.azure.com/mai/v1", api_version="preview", ) - assert ( - url - == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" - ) + assert url == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" def test_get_azure_ai_image_generation_config_returns_mai(self): config = get_azure_ai_image_generation_config("MAI-Image-2.5") @@ -124,13 +111,13 @@ class TestAzureMAIImageGeneration: config = AzureFoundryMAIImageGenerationConfig() optional_params = get_optional_params_image_gen( model="MAI-Image-2.5", - size="1792x1024", + size="1024x1024", n=1, custom_llm_provider="azure_ai", provider_config=config, drop_params=True, ) - assert optional_params["width"] == 1792 + assert optional_params["width"] == 1024 assert optional_params["height"] == 1024 assert "size" not in optional_params @@ -147,10 +134,7 @@ class TestAzureMAIImageGeneration: assert "api-version=preview" in url def test_mai_json_body_keeps_model(self): - api = ( - "https://my-resource.services.ai.azure.com/mai/v1/images/generations" - "?api-version=preview" - ) + api = "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" data = { "model": "MAI-Image-2.5", "prompt": "A photograph of a red fox", @@ -214,6 +198,74 @@ class TestAzureMAIImageGeneration: drop_params=True, ) + @pytest.mark.parametrize("size", ["512x512", "256x256", "700x1400"]) + def test_map_openai_params_size_below_minimum_dimension_raises(self, size): + """MAI requires >= 768px per side; the OpenAI size table offered smaller ones.""" + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(ValueError, match="at least 768 pixels"): + config.map_openai_params( + non_default_params={"size": size}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + + @pytest.mark.parametrize("size", ["1792x1024", "1024x1792"]) + def test_map_openai_params_size_over_total_pixel_budget_raises(self, size): + """MAI caps total pixels at 1024*1024, so both landscape/portrait sizes 400 upstream.""" + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(ValueError, match="at most 1048576 total pixels"): + config.map_openai_params( + non_default_params={"size": size}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + + def test_map_openai_params_explicit_width_height_not_range_checked(self): + """width/height pass through unmapped, so a future model's bounds stay reachable.""" + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"width": 1792, "height": 1024}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["width"] == 1792 + assert optional_params["height"] == 1024 + + @pytest.mark.parametrize("n", [2, 4]) + def test_map_openai_params_multi_image_n_raises(self, n): + """The MAI endpoint returns one image and ignores any count, so n>1 must not pass silently.""" + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(ValueError, match="returns exactly 1 image per request"): + config.map_openai_params( + non_default_params={"n": n}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + + def test_map_openai_params_multi_image_n_dropped_with_drop_params(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"n": 4}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert "n" not in optional_params + + def test_map_openai_params_single_image_n_still_passes_through(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"n": 1}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + assert optional_params["n"] == 1 + def test_map_openai_params_unsupported_param_raises(self): config = AzureFoundryMAIImageGenerationConfig() with pytest.raises(ValueError, match="Parameter quality is not supported"): @@ -363,16 +415,12 @@ class TestAzureMAIImageGeneration: litellm.model_cost = litellm.get_model_cost_map(url="") model = "azure_ai/MAI-Image-2.5" model_info = litellm.get_model_info(model=model, custom_llm_provider="azure_ai") - image_response = ImageResponse( - data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")] - ) + image_response = ImageResponse(data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]) cost = azure_ai_image_cost_calculator( model=model, image_response=image_response, ) - assert ( - cost == len(image_response.data or []) * model_info["output_cost_per_image"] - ) + assert cost == len(image_response.data or []) * model_info["output_cost_per_image"] assert cost > 0