feat(aiml): add openai/gpt-image-2 image model (#31323)

* feat(aiml): add openai/gpt-image-2 image model

Adds aiml/openai/gpt-image-2 to the cost map and teaches AimlImageGenerationConfig
to route OpenAI-style image models through the upstream OpenAI request schema
instead of the AI/ML flux schema. Without this, size, n, and response_format would
be remapped to image_size/num_images/output_format, which the gpt-image-2 endpoint
on api.aimlapi.com does not accept.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* chore(aiml): note gpt-image-2 flat-rate pricing basis; apply ruff format

Documents in the cost-map notes that output_cost_per_image is AI/ML's
published medium-quality rate, billed as a flat per-image price like the
other aiml image entries. Reformats the touched files under the repo's
ruff formatter (migrated from black in #31317).

* fix(aiml): drop /v1/images/edits from gpt-image-2 supported_endpoints

LiteLLM only implements an image generation transformer for AIML, so
listing /v1/images/edits overclaimed support. Align with every other
aiml image entry, which lists only /v1/images/generations.

* style(aiml): format transformation.py at line-length 88

The repo formats litellm/ with ruff at line-length 88 (Makefile/CI call
sites), while ruff.toml's global 120 only governs E501/import sorting.
Reformat the transformer to 88 so make format-check / CI lint pass, and
restore the test files to their original layout since tests/ is not part
of the auto-formatted tree.

---------

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

litellm/llms/aiml/image_generation/transformation.py

@@ -21,16 +21,39 @@ else:
     LiteLLMLoggingObj = Any
 
 
+OPENAI_STYLE_IMAGE_MODEL_PREFIXES: tuple[str, ...] = ("openai/",)
+
+
 class AimlImageGenerationConfig(BaseImageGenerationConfig):
     DEFAULT_BASE_URL: str = "https://api.aimlapi.com"
     IMAGE_GENERATION_ENDPOINT: str = "v1/images/generations"
 
+    @staticmethod
+    def _is_openai_style_model(model: str) -> bool:
+        """
+        OpenAI image models routed through AI/ML API (e.g. ``openai/gpt-image-2``)
+        use the upstream OpenAI request schema, not the flux-style schema used by
+        the rest of the AI/ML catalog.
+        """
+        return model.startswith(OPENAI_STYLE_IMAGE_MODEL_PREFIXES)
+
     def get_supported_openai_params(
         self, model: str
     ) -> List[OpenAIImageGenerationOptionalParams]:
         """
         https://api.aimlapi.com/v1/images/generations
         """
+        if self._is_openai_style_model(model):
+            return [
+                "n",
+                "size",
+                "quality",
+                "response_format",
+                "output_format",
+                "background",
+                "moderation",
+                "output_compression",
+            ]
         return ["n", "response_format", "size"]
 
     def map_openai_params(
@@ -41,39 +64,38 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig):
         drop_params: bool,
     ) -> dict:
         supported_params = self.get_supported_openai_params(model)
+        is_openai_style = self._is_openai_style_model(model)
 
         for k in non_default_params.keys():
-            if k not in optional_params.keys():
-                if k in supported_params:
-                    # Map OpenAI params to AI/ML params
-                    if k == "n":
-                        optional_params["num_images"] = non_default_params[k]
-                    elif k == "response_format":
-                        optional_params["output_format"] = non_default_params[k]
-                    elif k == "size":
-                        # Map OpenAI size format to AI/ML image_size
-                        size_value = non_default_params[k]
-                        if isinstance(size_value, str):
-                            # Handle standard OpenAI sizes like "1024x1024"
-                            if "x" in size_value:
-                                width, height = map(int, size_value.split("x"))
-                                optional_params["image_size"] = {
-                                    "width": width,
-                                    "height": height,
-                                }
-                            else:
-                                # Pass through predefined sizes
-                                optional_params["image_size"] = size_value
-                        else:
-                            optional_params["image_size"] = size_value
-                    else:
-                        optional_params[k] = non_default_params[k]
-                elif drop_params:
-                    pass
+            if k in optional_params.keys():
+                continue
+            if k not in supported_params:
+                if drop_params:
+                    continue
+                raise ValueError(
+                    f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters."
+                )
+
+            if is_openai_style:
+                optional_params[k] = non_default_params[k]
+                continue
+
+            if k == "n":
+                optional_params["num_images"] = non_default_params[k]
+            elif k == "response_format":
+                optional_params["output_format"] = non_default_params[k]
+            elif k == "size":
+                size_value = non_default_params[k]
+                if isinstance(size_value, str) and "x" in size_value:
+                    width, height = map(int, size_value.split("x"))
+                    optional_params["image_size"] = {
+                        "width": width,
+                        "height": height,
+                    }
                 else:
-                    raise ValueError(
-                        f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters."
-                    )
+                    optional_params["image_size"] = size_value
+            else:
+                optional_params[k] = non_default_params[k]
 
         return optional_params
 
@@ -131,10 +153,13 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig):
         headers: dict,
     ) -> dict:
         """
-        Transform the image generation request to the AI/ML flux image generation request body
+        Transform the image generation request to the AI/ML image generation request body
 
         https://api.aimlapi.com/v1/images/generations
         """
+        if self._is_openai_style_model(model):
+            return {"model": model, "prompt": prompt, **optional_params}
+
         aiml_image_generation_request_body: AimlImageGenerationRequestParams = (
             AimlImageGenerationRequestParams(
                 prompt=prompt,

litellm/model_prices_and_context_window_backup.json

@@ -273,6 +273,19 @@
             "/v1/images/generations"
         ]
     },
+    "aiml/openai/gpt-image-2": {
+        "litellm_provider": "aiml",
+        "metadata": {
+            "notes": "OpenAI gpt-image-2 via AI/ML API - flagship multimodal image generation and editing model with reasoning and 2K output. output_cost_per_image is AI/ML's published medium-quality rate; like the other aiml image entries it is billed as a flat per-image price"
+        },
+        "mode": "image_generation",
+        "output_cost_per_image": 0.054,
+        "source": "https://docs.aimlapi.com/api-references/image-models/openai/gpt-image-2",
+        "supported_endpoints": [
+            "/v1/images/generations"
+        ],
+        "supports_vision": true
+    },
     "amazon.nova-canvas-v1:0": {
         "litellm_provider": "bedrock",
         "max_input_tokens": 2600,

model_prices_and_context_window.json

@@ -273,6 +273,19 @@
             "/v1/images/generations"
         ]
     },
+    "aiml/openai/gpt-image-2": {
+        "litellm_provider": "aiml",
+        "metadata": {
+            "notes": "OpenAI gpt-image-2 via AI/ML API - flagship multimodal image generation and editing model with reasoning and 2K output. output_cost_per_image is AI/ML's published medium-quality rate; like the other aiml image entries it is billed as a flat per-image price"
+        },
+        "mode": "image_generation",
+        "output_cost_per_image": 0.054,
+        "source": "https://docs.aimlapi.com/api-references/image-models/openai/gpt-image-2",
+        "supported_endpoints": [
+            "/v1/images/generations"
+        ],
+        "supports_vision": true
+    },
     "amazon.nova-canvas-v1:0": {
         "litellm_provider": "bedrock",
         "max_input_tokens": 2600,

tests/image_gen_tests/test_image_generation.py

@@ -386,6 +386,62 @@ async def test_aiml_image_generation_with_dynamic_api_key():
         assert captured_json_data["model"] == "flux-pro/v1.1"
 
 
+@pytest.mark.asyncio
+async def test_aiml_openai_gpt_image_2_request_uses_openai_param_shape():
+    """End-to-end check that ``aiml/openai/gpt-image-2`` keeps the upstream
+    OpenAI request shape (``size``/``n``/``response_format``) instead of
+    being remapped to the AI/ML flux schema (``image_size``/``num_images``/
+    ``output_format``), and hits the correct upstream model name.
+    """
+    from unittest.mock import MagicMock, patch
+    import json as _json
+
+    mock_aiml_response = {
+        "created": 1703658209,
+        "data": [{"url": "https://example.com/gpt-image-2.png"}],
+    }
+
+    captured = {}
+
+    def capture_post_call(*args, **kwargs):
+        captured["url"] = kwargs.get("url") or (args[0] if args else None)
+        captured["headers"] = kwargs.get("headers", {})
+        captured["json"] = kwargs.get("json", {})
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = mock_aiml_response
+        mock_response.text = _json.dumps(mock_aiml_response)
+        return mock_response
+
+    with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post:
+        mock_post.side_effect = capture_post_call
+
+        await litellm.aimage_generation(
+            prompt="A T-Rex relaxing on a beach",
+            model="aiml/openai/gpt-image-2",
+            api_key="test-key-mocked-no-credits-needed",
+            size="1024x1536",
+            quality="high",
+            response_format="b64_json",
+            n=1,
+        )
+
+    assert captured["url"] is not None
+    assert "api.aimlapi.com" in captured["url"]
+    assert "/v1/images/generations" in captured["url"]
+
+    body = captured["json"]
+    assert body["model"] == "openai/gpt-image-2"
+    assert body["prompt"] == "A T-Rex relaxing on a beach"
+    assert body["size"] == "1024x1536"
+    assert body["quality"] == "high"
+    assert body["response_format"] == "b64_json"
+    assert body["n"] == 1
+    assert "image_size" not in body
+    assert "num_images" not in body
+    assert "output_format" not in body
+
+
 @pytest.mark.asyncio
 async def test_azure_image_generation_request_body():
     """Azure deployment URL selects the model; JSON body omits ``model`` (#26316)."""

tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py

@@ -0,0 +1,147 @@
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../../../.."))
+
+os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
+
+import litellm
+
+litellm.model_cost = litellm.get_model_cost_map(url="")
+
+from litellm.llms.aiml.image_generation.cost_calculator import (
+    cost_calculator as aiml_cost_calculator,
+)
+from litellm.llms.aiml.image_generation.transformation import (
+    AimlImageGenerationConfig,
+)
+from litellm.types.utils import ImageObject, ImageResponse
+
+
+def test_openai_style_model_supports_full_openai_param_surface():
+    params = AimlImageGenerationConfig().get_supported_openai_params(
+        "openai/gpt-image-2"
+    )
+    assert {
+        "n",
+        "size",
+        "quality",
+        "response_format",
+        "output_format",
+        "background",
+        "moderation",
+        "output_compression",
+    } == set(params)
+
+
+def test_flux_style_model_keeps_legacy_param_surface():
+    assert AimlImageGenerationConfig().get_supported_openai_params("flux-pro/v1.1") == [
+        "n",
+        "response_format",
+        "size",
+    ]
+
+
+def test_openai_style_request_passes_params_through_unchanged():
+    """gpt-image-2 must receive OpenAI-shaped fields (size string, n, response_format) verbatim;
+    the flux-style remapping to ``num_images``/``image_size``/``output_format`` would break the upstream call.
+    """
+    config = AimlImageGenerationConfig()
+    mapped = config.map_openai_params(
+        non_default_params={
+            "n": 1,
+            "size": "1024x1536",
+            "quality": "high",
+            "response_format": "b64_json",
+            "output_format": "png",
+        },
+        optional_params={},
+        model="openai/gpt-image-2",
+        drop_params=False,
+    )
+
+    body = config.transform_image_generation_request(
+        model="openai/gpt-image-2",
+        prompt="A cute baby sea otter",
+        optional_params=mapped,
+        litellm_params={},
+        headers={},
+    )
+
+    assert body == {
+        "model": "openai/gpt-image-2",
+        "prompt": "A cute baby sea otter",
+        "n": 1,
+        "size": "1024x1536",
+        "quality": "high",
+        "response_format": "b64_json",
+        "output_format": "png",
+    }
+
+
+def test_flux_style_request_still_remaps_to_legacy_fields():
+    config = AimlImageGenerationConfig()
+    mapped = config.map_openai_params(
+        non_default_params={
+            "n": 2,
+            "size": "1024x1024",
+            "response_format": "png",
+        },
+        optional_params={},
+        model="flux-pro/v1.1",
+        drop_params=False,
+    )
+
+    body = config.transform_image_generation_request(
+        model="flux-pro/v1.1",
+        prompt="hello",
+        optional_params=mapped,
+        litellm_params={},
+        headers={},
+    )
+
+    assert body["model"] == "flux-pro/v1.1"
+    assert body["prompt"] == "hello"
+    assert body["num_images"] == 2
+    assert body["image_size"] == {"width": 1024, "height": 1024}
+    assert body["output_format"] == "png"
+    assert "n" not in body
+    assert "size" not in body
+    assert "response_format" not in body
+
+
+def test_openai_style_unsupported_param_raises_without_drop_params():
+    with pytest.raises(ValueError):
+        AimlImageGenerationConfig().map_openai_params(
+            non_default_params={"image_size": {"width": 1024, "height": 1024}},
+            optional_params={},
+            model="openai/gpt-image-2",
+            drop_params=False,
+        )
+
+
+def test_openai_style_unsupported_param_dropped_with_drop_params():
+    mapped = AimlImageGenerationConfig().map_openai_params(
+        non_default_params={"image_size": {"width": 1024, "height": 1024}},
+        optional_params={},
+        model="openai/gpt-image-2",
+        drop_params=True,
+    )
+    assert mapped == {}
+
+
+def test_cost_calculator_uses_aiml_pricing_for_gpt_image_2():
+    """Regression: pricing must come from the ``aiml/openai/gpt-image-2`` entry,
+    not the upstream OpenAI token-based entry.
+    """
+    response = ImageResponse(
+        data=[
+            ImageObject(b64_json=None, url="https://example.com/1.png"),
+            ImageObject(b64_json=None, url="https://example.com/2.png"),
+        ]
+    )
+    assert aiml_cost_calculator(
+        model="openai/gpt-image-2", image_response=response
+    ) == pytest.approx(0.054 * 2)