From b7f28bd89f89b8efb54c9ed6aefa80216aa1e44f Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 25 Jun 2026 16:41:43 -0700 Subject: [PATCH] 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 * 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 Co-authored-by: Mateo Wang --- .../aiml/image_generation/transformation.py | 87 +++++++---- ...odel_prices_and_context_window_backup.json | 13 ++ model_prices_and_context_window.json | 13 ++ .../image_gen_tests/test_image_generation.py | 56 +++++++ ...st_aiml_image_generation_transformation.py | 147 ++++++++++++++++++ 5 files changed, 285 insertions(+), 31 deletions(-) create mode 100644 tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py diff --git a/litellm/llms/aiml/image_generation/transformation.py b/litellm/llms/aiml/image_generation/transformation.py index 39b1cc742d4..92a1510f3a1 100644 --- a/litellm/llms/aiml/image_generation/transformation.py +++ b/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, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6ebac7efc8d..296996d186e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/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, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4ec4d8bb884..64f23118275 100644 --- a/model_prices_and_context_window.json +++ b/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, diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index 873777189c9..f0a73325afa 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/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).""" diff --git a/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py b/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py new file mode 100644 index 00000000000..7485f2121df --- /dev/null +++ b/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)