From bc29a10c6cde7ec7075f74d3af4b608853f70034 Mon Sep 17 00:00:00 2001 From: Dmitriy Alergant Date: Mon, 18 May 2026 16:47:46 -0400 Subject: [PATCH 01/12] Fix Gemini image config mapping --- litellm/llms/gemini/common_utils.py | 146 +++++++++++++ .../llms/gemini/image_edit/transformation.py | 43 ++-- .../gemini/image_generation/transformation.py | 76 ++++--- litellm/types/images/main.py | 1 + litellm/types/llms/openai.py | 1 + litellm/utils.py | 2 + tests/llm_translation/test_gemini.py | 200 ++++++++++++++++++ tests/proxy_unit_tests/test_proxy_server.py | 2 + .../test_gemini_image_edit_transformation.py | 66 +++++- 9 files changed, 487 insertions(+), 50 deletions(-) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index bc963d62b5f5..0c7bcdb45b37 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -1,5 +1,6 @@ import base64 import datetime +import math from typing import Any, Dict, List, Optional, Union import httpx @@ -13,6 +14,151 @@ from litellm.types.utils import TokenCountResponse +GEMINI_IMAGE_ASPECT_RATIOS: Dict[str, float] = { + "1:1": 1 / 1, + "1:4": 1 / 4, + "1:8": 1 / 8, + "2:3": 2 / 3, + "3:2": 3 / 2, + "3:4": 3 / 4, + "4:1": 4 / 1, + "4:3": 4 / 3, + "4:5": 4 / 5, + "5:4": 5 / 4, + "8:1": 8 / 1, + "9:16": 9 / 16, + "16:9": 16 / 9, + "21:9": 21 / 9, +} + +GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO: Dict[tuple[int, int], str] = { + (512, 512): "1:1", + (1024, 1024): "1:1", + (2048, 2048): "1:1", + (4096, 4096): "1:1", + (256, 1024): "1:4", + (512, 2048): "1:4", + (1024, 4096): "1:4", + (2048, 8192): "1:4", + (192, 1536): "1:8", + (384, 3072): "1:8", + (768, 6144): "1:8", + (1536, 12288): "1:8", + (424, 632): "2:3", + (848, 1264): "2:3", + (1696, 2528): "2:3", + (3392, 5056): "2:3", + (632, 424): "3:2", + (1264, 848): "3:2", + (2528, 1696): "3:2", + (5056, 3392): "3:2", + (448, 600): "3:4", + (896, 1200): "3:4", + (1792, 2400): "3:4", + (3584, 4800): "3:4", + (1024, 256): "4:1", + (2048, 512): "4:1", + (4096, 1024): "4:1", + (8192, 2048): "4:1", + (600, 448): "4:3", + (1200, 896): "4:3", + (2400, 1792): "4:3", + (4800, 3584): "4:3", + (464, 576): "4:5", + (928, 1152): "4:5", + (1856, 2304): "4:5", + (3712, 4608): "4:5", + (576, 464): "5:4", + (1152, 928): "5:4", + (2304, 1856): "5:4", + (4608, 3712): "5:4", + (1536, 192): "8:1", + (3072, 384): "8:1", + (6144, 768): "8:1", + (12288, 1536): "8:1", + (384, 688): "9:16", + (768, 1376): "9:16", + (1536, 2752): "9:16", + (3072, 5504): "9:16", + (688, 384): "16:9", + (1376, 768): "16:9", + (2752, 1536): "16:9", + (5504, 3072): "16:9", + (792, 168): "21:9", + (1584, 672): "21:9", + (3168, 1344): "21:9", + (6336, 2688): "21:9", + (1280, 896): "4:3", + (896, 1280): "3:4", +} + + +def map_openai_size_to_gemini_image_config( + size: str, model: str +) -> Optional[Dict[str, str]]: + dimensions = _parse_openai_image_size(size) + if dimensions is None: + return None + + width, height = dimensions + image_config = { + "aspectRatio": _map_dimensions_to_gemini_aspect_ratio(width, height) + } + if supports_gemini_image_size(model): + image_config["imageSize"] = _map_dimensions_to_gemini_image_size(width, height) + return image_config + + +def supports_gemini_image_size(model: str) -> bool: + # Gemini 2.5 Flash image supports aspectRatio but rejects imageSize; newer + # Gemini image models are expected to support both fields. + return "2.5-flash" not in model + + +def _parse_openai_image_size(size: str) -> Optional[tuple[int, int]]: + if size == "auto": + return None + + width_str, separator, height_str = size.lower().partition("x") + if not separator: + return None + + try: + width = int(width_str) + height = int(height_str) + except ValueError: + return None + + if width <= 0 or height <= 0: + return None + + return width, height + + +def _map_dimensions_to_gemini_aspect_ratio(width: int, height: int) -> str: + if (width, height) in GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO: + return GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO[(width, height)] + + requested_ratio = width / height + return min( + GEMINI_IMAGE_ASPECT_RATIOS, + key=lambda aspect_ratio: abs( + math.log(GEMINI_IMAGE_ASPECT_RATIOS[aspect_ratio] / requested_ratio) + ), + ) + + +def _map_dimensions_to_gemini_image_size(width: int, height: int) -> str: + effective_square_side = math.sqrt(width * height) + if effective_square_side < 768: + return "512" + if effective_square_side < 1536: + return "1K" + if effective_square_side < 3072: + return "2K" + return "4K" + + class GeminiError(BaseLLMException): pass diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index c8aaab0e14ec..76c883fdfc6e 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -1,4 +1,5 @@ import base64 +import json from io import BufferedReader, BytesIO from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast @@ -7,6 +8,10 @@ from litellm.images.utils import ImageEditRequestUtils from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.llms.gemini.common_utils import ( + map_openai_size_to_gemini_image_config, + supports_gemini_image_size, +) from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.router import GenericLiteLLMParams @@ -22,7 +27,7 @@ class GeminiImageEditConfig(BaseImageEditConfig): DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta" - SUPPORTED_PARAMS: List[str] = ["size"] + SUPPORTED_PARAMS: List[str] = ["size", "imageConfig"] def get_supported_openai_params(self, model: str) -> List[str]: return list(self.SUPPORTED_PARAMS) @@ -43,9 +48,18 @@ def map_openai_params( mapped_params: Dict[str, Any] = {} if "size" in filtered_params: - mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio( - filtered_params["size"] # type: ignore[arg-type] + image_config = map_openai_size_to_gemini_image_config( + filtered_params["size"], # type: ignore[arg-type] + model, ) + if image_config is not None: + mapped_params["imageConfig"] = image_config + + image_config_param = filtered_params.get("imageConfig") + if isinstance(image_config_param, str): + image_config_param = json.loads(image_config_param) + if isinstance(image_config_param, dict): + mapped_params["imageConfig"] = image_config_param return mapped_params @@ -109,13 +123,12 @@ def transform_image_edit_request( # type: ignore[override] generation_config: Dict[str, Any] = {} - if "aspectRatio" in image_edit_optional_request_params: - # Move aspectRatio into imageConfig inside generationConfig - if "imageConfig" not in generation_config: - generation_config["imageConfig"] = {} - generation_config["imageConfig"]["aspectRatio"] = ( - image_edit_optional_request_params["aspectRatio"] - ) + if isinstance(image_edit_optional_request_params.get("imageConfig"), dict): + image_config = dict(image_edit_optional_request_params["imageConfig"]) + if not supports_gemini_image_size(model): + image_config.pop("imageSize", None) + if image_config: + generation_config["imageConfig"] = image_config if generation_config: request_body["generationConfig"] = generation_config @@ -158,16 +171,6 @@ def transform_image_edit_response( model_response.data = cast(List[OpenAIImage], data_list) return model_response - def _map_size_to_aspect_ratio(self, size: str) -> str: - aspect_ratio_map = { - "1024x1024": "1:1", - "1792x1024": "16:9", - "1024x1792": "9:16", - "1280x896": "4:3", - "896x1280": "3:4", - } - return aspect_ratio_map.get(size, "1:1") - def _prepare_inline_image_parts( self, image: Union[FileTypes, List[FileTypes]] ) -> List[Dict[str, Any]]: diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 9c4cd008b8cb..aa899bdbe8c9 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -1,10 +1,14 @@ -from typing import TYPE_CHECKING, Any, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional import httpx from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) +from litellm.llms.gemini.common_utils import ( + map_openai_size_to_gemini_image_config, + supports_gemini_image_size, +) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.gemini import GeminiImageGenerationRequest from litellm.types.llms.openai import ( @@ -36,7 +40,10 @@ def get_supported_openai_params( Google AI Imagen API supported parameters https://ai.google.dev/gemini-api/docs/imagen """ - return ["n", "size"] + supported_params = ["n", "size"] + if "gemini" in model: + supported_params.append("imageConfig") + return supported_params # type: ignore[return-value] def map_openai_params( self, @@ -48,32 +55,33 @@ def map_openai_params( supported_params = self.get_supported_openai_params(model) mapped_params = {} - for k, v in non_default_params.items(): - if k not in optional_params.keys(): - if k in supported_params: - # Map OpenAI parameters to Google format - if k == "n": - mapped_params["sampleCount"] = v - elif k == "size": - # Map OpenAI size format to Google aspectRatio - mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(v) - else: - mapped_params[k] = v - return mapped_params + if "n" in non_default_params and "n" not in optional_params: + mapped_params["sampleCount"] = non_default_params["n"] - def _map_size_to_aspect_ratio(self, size: str) -> str: - """ - https://ai.google.dev/gemini-api/docs/image-generation + if "size" in non_default_params and "size" not in optional_params: + image_config = map_openai_size_to_gemini_image_config( + non_default_params["size"], model + ) + if image_config is not None: + if "gemini" in model: + mapped_params["imageConfig"] = image_config + else: + mapped_params["aspectRatio"] = image_config["aspectRatio"] - """ - aspect_ratio_map = { - "1024x1024": "1:1", - "1792x1024": "16:9", - "1024x1792": "9:16", - "1280x896": "4:3", - "896x1280": "3:4", - } - return aspect_ratio_map.get(size, "1:1") + if ( + "imageConfig" in supported_params + and isinstance(non_default_params.get("imageConfig"), dict) + ): + mapped_params["imageConfig"] = non_default_params["imageConfig"] + + for k, v in non_default_params.items(): + if ( + k not in ("n", "size", "imageConfig") + and k not in optional_params + and k in supported_params + ): + mapped_params[k] = v + return mapped_params def _transform_image_usage(self, usage_metadata: dict) -> ImageUsage: """ @@ -180,9 +188,23 @@ def transform_image_generation_request( """ # For Gemini Flash Image Preview models, use standard Gemini format if "gemini" in model: + generation_config: Dict[str, Any] = { + "response_modalities": ["IMAGE", "TEXT"] + } + image_config: Dict[str, Any] = {} + + if isinstance(optional_params.get("imageConfig"), dict): + image_config.update(optional_params["imageConfig"]) + + if not supports_gemini_image_size(model): + image_config.pop("imageSize", None) + + if image_config: + generation_config["imageConfig"] = image_config + request_body: dict = { "contents": [{"parts": [{"text": prompt}]}], - "generationConfig": {"response_modalities": ["IMAGE", "TEXT"]}, + "generationConfig": generation_config, } return request_body else: diff --git a/litellm/types/images/main.py b/litellm/types/images/main.py index 819f49545897..80e55297c423 100644 --- a/litellm/types/images/main.py +++ b/litellm/types/images/main.py @@ -20,6 +20,7 @@ class ImageEditOptionalRequestParams(TypedDict, total=False): response_format: Optional[Literal["url", "b64_json"]] size: Optional[str] user: Optional[str] + imageConfig: Optional[Dict[str, Any]] class ImageEditRequestParams(ImageEditOptionalRequestParams, total=False): diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index abe58199dfdc..0dbf35e149cf 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1076,6 +1076,7 @@ class LiteLLMFineTuningJobCreate(FineTuningJobCreate): "image_url", "image_prompt_strength", "aspect_ratio", + "imageConfig", ] OpenAIImageEditOptionalParams = Literal[ diff --git a/litellm/utils.py b/litellm/utils.py index 001c89fee4c6..76bf15f96f26 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3140,6 +3140,7 @@ def get_optional_params_image_gen( size: Optional[str] = None, style: Optional[str] = None, user: Optional[str] = None, + imageConfig: Optional[dict] = None, custom_llm_provider: Optional[str] = None, additional_drop_params: Optional[list] = None, provider_config: Optional[BaseImageGenerationConfig] = None, @@ -3176,6 +3177,7 @@ def get_optional_params_image_gen( "size": None, "style": None, "user": None, + "imageConfig": None, } non_default_params = _get_non_default_params( diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index bd340aa63be2..b3463863ba7c 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -17,6 +17,66 @@ import json +GEMINI_3_IMAGE_SIZE_MAPPINGS = [ + ("512x512", "1:1", "512"), + ("1024x1024", "1:1", "1K"), + ("2048x2048", "1:1", "2K"), + ("4096x4096", "1:1", "4K"), + ("256x1024", "1:4", "512"), + ("512x2048", "1:4", "1K"), + ("1024x4096", "1:4", "2K"), + ("2048x8192", "1:4", "4K"), + ("192x1536", "1:8", "512"), + ("384x3072", "1:8", "1K"), + ("768x6144", "1:8", "2K"), + ("1536x12288", "1:8", "4K"), + ("424x632", "2:3", "512"), + ("848x1264", "2:3", "1K"), + ("1696x2528", "2:3", "2K"), + ("3392x5056", "2:3", "4K"), + ("632x424", "3:2", "512"), + ("1264x848", "3:2", "1K"), + ("2528x1696", "3:2", "2K"), + ("5056x3392", "3:2", "4K"), + ("448x600", "3:4", "512"), + ("896x1200", "3:4", "1K"), + ("1792x2400", "3:4", "2K"), + ("3584x4800", "3:4", "4K"), + ("1024x256", "4:1", "512"), + ("2048x512", "4:1", "1K"), + ("4096x1024", "4:1", "2K"), + ("8192x2048", "4:1", "4K"), + ("600x448", "4:3", "512"), + ("1200x896", "4:3", "1K"), + ("2400x1792", "4:3", "2K"), + ("4800x3584", "4:3", "4K"), + ("464x576", "4:5", "512"), + ("928x1152", "4:5", "1K"), + ("1856x2304", "4:5", "2K"), + ("3712x4608", "4:5", "4K"), + ("576x464", "5:4", "512"), + ("1152x928", "5:4", "1K"), + ("2304x1856", "5:4", "2K"), + ("4608x3712", "5:4", "4K"), + ("1536x192", "8:1", "512"), + ("3072x384", "8:1", "1K"), + ("6144x768", "8:1", "2K"), + ("12288x1536", "8:1", "4K"), + ("384x688", "9:16", "512"), + ("768x1376", "9:16", "1K"), + ("1536x2752", "9:16", "2K"), + ("3072x5504", "9:16", "4K"), + ("688x384", "16:9", "512"), + ("1376x768", "16:9", "1K"), + ("2752x1536", "16:9", "2K"), + ("5504x3072", "16:9", "4K"), + ("792x168", "21:9", "512"), + ("1584x672", "21:9", "1K"), + ("3168x1344", "21:9", "2K"), + ("6336x2688", "21:9", "4K"), +] + + class TestGoogleAIStudioGemini(BaseLLMChatTest): def get_base_completion_call_args(self) -> dict: return {"model": "gemini/gemini-2.5-flash"} @@ -365,6 +425,143 @@ def test_gemini_flash_image_preview_models(model_name: str): ] +@pytest.mark.parametrize( + "model, kwargs, expected_image_config", + [ + ( + "gemini/gemini-3-pro-image-preview", + {"imageConfig": {"aspectRatio": "16:9", "imageSize": "512px"}}, + {"aspectRatio": "16:9", "imageSize": "512px"}, + ), + ( + "gemini/gemini-2.5-flash-image", + {"size": "2048x2048"}, + {"aspectRatio": "1:1"}, + ), + ], +) +def test_gemini_image_generation_forwards_image_config( + model: str, kwargs: dict, expected_image_config: dict +): + from unittest.mock import patch, MagicMock + + with patch( + "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post" + ) as mock_post: + mock_http_response = MagicMock() + mock_http_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [{"inlineData": {"data": "test_base64_image_data"}}] + } + } + ] + } + mock_http_response.status_code = 200 + mock_post.return_value = mock_http_response + + litellm.image_generation( + model=model, + prompt="Generate a simple test image", + api_key="test_api_key", + **kwargs, + ) + + request_data = mock_post.call_args.kwargs.get("json", {}) + assert request_data["generationConfig"]["imageConfig"] == expected_image_config + + +def test_gemini_image_generation_image_config_takes_precedence_over_size(): + from litellm.llms.gemini.image_generation.transformation import GoogleImageGenConfig + + explicit_image_config = {"aspectRatio": "16:9", "imageSize": "2K"} + + mapped_params = GoogleImageGenConfig().map_openai_params( + non_default_params={ + "imageConfig": explicit_image_config, + "size": "768x1376", + }, + optional_params={}, + model="gemini-3-pro-image-preview", + drop_params=False, + ) + + assert mapped_params["imageConfig"] == explicit_image_config + + +def test_gemini_image_generation_ignores_non_dict_image_config(): + from litellm.llms.gemini.image_generation.transformation import GoogleImageGenConfig + + mapped_params = GoogleImageGenConfig().map_openai_params( + non_default_params={ + "size": "768x1376", + "imageConfig": "not-a-dict", + }, + optional_params={}, + model="gemini-3-pro-image-preview", + drop_params=False, + ) + + assert mapped_params["imageConfig"] == {"aspectRatio": "9:16", "imageSize": "1K"} + + +@pytest.mark.parametrize( + "size, expected_aspect_ratio, expected_image_size", + GEMINI_3_IMAGE_SIZE_MAPPINGS, +) +def test_gemini_image_generation_openai_size_maps_to_google_table( + size: str, expected_aspect_ratio: str, expected_image_size: str +): + from litellm.llms.gemini.common_utils import ( + map_openai_size_to_gemini_image_config, + ) + + assert map_openai_size_to_gemini_image_config( + size, "gemini-3-pro-image-preview" + ) == { + "aspectRatio": expected_aspect_ratio, + "imageSize": expected_image_size, + } + + +@pytest.mark.parametrize( + "size, expected_aspect_ratio, expected_image_size", + [ + ("1000x1800", "9:16", "1K"), + ("1800x1000", "16:9", "1K"), + ("3000x3000", "1:1", "2K"), + ("500x500", "1:1", "512"), + ("1280x896", "4:3", "1K"), + ("896x1280", "3:4", "1K"), + ], +) +def test_gemini_image_generation_openai_size_snaps_to_nearest_option( + size: str, expected_aspect_ratio: str, expected_image_size: str +): + from litellm.llms.gemini.common_utils import ( + map_openai_size_to_gemini_image_config, + ) + + assert map_openai_size_to_gemini_image_config( + size, "gemini-3-pro-image-preview" + ) == { + "aspectRatio": expected_aspect_ratio, + "imageSize": expected_image_size, + } + + +@pytest.mark.parametrize("size", ["auto", "invalid", "0x1024", "1024x0"]) +def test_gemini_image_generation_openai_size_auto_uses_google_defaults(size: str): + from litellm.llms.gemini.common_utils import ( + map_openai_size_to_gemini_image_config, + ) + + assert map_openai_size_to_gemini_image_config( + size, "gemini-3-pro-image-preview" + ) is None + + def test_gemini_imagen_models_use_predict_endpoint(): """ Test that Imagen models still use :predict endpoint (not broken by gemini-2.5-flash-image-preview fix) @@ -387,6 +584,7 @@ def test_gemini_imagen_models_use_predict_endpoint(): response = litellm.image_generation( model="gemini/imagen-3.0-generate-001", prompt="Generate a simple test image", + size="1280x896", api_key="test_api_key", ) @@ -410,6 +608,8 @@ def test_gemini_imagen_models_use_predict_endpoint(): request_data = call_args.kwargs.get("json", {}) assert "instances" in request_data assert "parameters" in request_data + assert request_data["parameters"]["aspectRatio"] == "4:3" + assert "imageConfig" not in request_data["parameters"] def test_gemini_thinking(): diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 9c08175767de..6fdabe64e251 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -804,6 +804,7 @@ def test_img_gen(mock_aimage_generation, client_no_auth): "prompt": "A cute baby sea otter", "n": 1, "size": "1024x1024", + "imageConfig": {"aspectRatio": "9:16", "imageSize": "1K"}, } response = client_no_auth.post("/v1/images/generations", json=test_data) @@ -813,6 +814,7 @@ def test_img_gen(mock_aimage_generation, client_no_auth): prompt="A cute baby sea otter", n=1, size="1024x1024", + imageConfig={"aspectRatio": "9:16", "imageSize": "1K"}, metadata=mock.ANY, proxy_server_request=mock.ANY, secret_fields=mock.ANY, diff --git a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py index 682df923693d..71afcc9a3b45 100644 --- a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py +++ b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py @@ -30,20 +30,60 @@ def test_map_openai_params(self) -> None: drop_params=False, ) - assert mapped["aspectRatio"] == "16:9" + assert mapped["imageConfig"] == {"aspectRatio": "16:9"} assert "response_format" not in mapped assert "quality" not in mapped + def test_map_openai_params_with_image_size_for_gemini_3(self) -> None: + optional_params: Dict[str, object] = { + "size": "768x1376", + } + + mapped = self.config.map_openai_params( + image_edit_optional_params=optional_params, # type: ignore[arg-type] + model="gemini-3-pro-image-preview", + drop_params=False, + ) + + assert mapped["imageConfig"] == {"aspectRatio": "9:16", "imageSize": "1K"} + + def test_map_openai_params_forwards_image_config_as_is(self) -> None: + optional_params: Dict[str, object] = { + "size": "1024x1024", + "imageConfig": {"aspectRatio": "16:9", "imageSize": "512px"}, + } + + mapped = self.config.map_openai_params( + image_edit_optional_params=optional_params, # type: ignore[arg-type] + model="gemini-3-pro-image-preview", + drop_params=False, + ) + + assert mapped["imageConfig"] == {"aspectRatio": "16:9", "imageSize": "512px"} + + def test_map_openai_params_parses_form_image_config_json(self) -> None: + optional_params: Dict[str, object] = { + "imageConfig": '{"aspectRatio":"16:9","imageSize":"1K"}', + } + + mapped = self.config.map_openai_params( + image_edit_optional_params=optional_params, # type: ignore[arg-type] + model="gemini-3-pro-image-preview", + drop_params=False, + ) + + assert mapped["imageConfig"] == {"aspectRatio": "16:9", "imageSize": "1K"} + def test_transform_image_edit_request(self) -> None: image_bytes = b"fake_image_data" image = BytesIO(image_bytes) optional_params = { "sampleCount": 2, - "aspectRatio": "16:9", + "imageConfig": {"aspectRatio": "16:9", "imageSize": "2K"}, } request_body, files = self.config.transform_image_edit_request( - model=self.model, + model="gemini-3-pro-image-preview", prompt=self.prompt, image=[image], # Gemini pipeline passes list of images image_edit_optional_request_params=optional_params, @@ -62,6 +102,26 @@ def test_transform_image_edit_request(self) -> None: generation_config = request_body["generationConfig"] assert generation_config["imageConfig"]["aspectRatio"] == "16:9" + assert generation_config["imageConfig"]["imageSize"] == "2K" + + def test_transform_image_edit_request_omits_image_size_for_gemini_25(self) -> None: + image = BytesIO(b"fake_image_data") + optional_params = { + "imageConfig": {"aspectRatio": "16:9", "imageSize": "2K"}, + } + + request_body, _ = self.config.transform_image_edit_request( + model=self.model, + prompt=self.prompt, + image=[image], + image_edit_optional_request_params=optional_params, + litellm_params=MagicMock(), + headers={}, + ) + + assert request_body["generationConfig"]["imageConfig"] == { + "aspectRatio": "16:9" + } def test_transform_image_edit_request_multiple_images(self) -> None: image_one = BytesIO(b"image_one") From 0e77cc4525929d4ba330a37085a3fcd37101f365 Mon Sep 17 00:00:00 2001 From: Dmitriy Alergant Date: Mon, 18 May 2026 18:26:23 -0400 Subject: [PATCH 02/12] Address Gemini image config review --- litellm/llms/gemini/common_utils.py | 2 +- .../llms/gemini/image_edit/transformation.py | 9 ++++++++- .../gemini/image_generation/transformation.py | 2 ++ litellm/types/llms/gemini.py | 3 +++ tests/llm_translation/test_gemini.py | 3 ++- .../test_gemini_image_edit_transformation.py | 17 +++++++++++++++++ 6 files changed, 33 insertions(+), 3 deletions(-) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 0c7bcdb45b37..7c7c38d760d3 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -84,7 +84,7 @@ (1376, 768): "16:9", (2752, 1536): "16:9", (5504, 3072): "16:9", - (792, 168): "21:9", + (792, 336): "21:9", (1584, 672): "21:9", (3168, 1344): "21:9", (6336, 2688): "21:9", diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index 76c883fdfc6e..75115b2653b2 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -6,6 +6,7 @@ import httpx from httpx._types import RequestFiles +import litellm from litellm.images.utils import ImageEditRequestUtils from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.gemini.common_utils import ( @@ -57,7 +58,13 @@ def map_openai_params( image_config_param = filtered_params.get("imageConfig") if isinstance(image_config_param, str): - image_config_param = json.loads(image_config_param) + try: + image_config_param = json.loads(image_config_param) + except json.JSONDecodeError as exc: + raise litellm.UnsupportedParamsError( + model=model, + message="`imageConfig` must be valid JSON when provided as a string.", + ) from exc if isinstance(image_config_param, dict): mapped_params["imageConfig"] = image_config_param diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index aa899bdbe8c9..d012037bca3c 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -67,6 +67,8 @@ def map_openai_params( mapped_params["imageConfig"] = image_config else: mapped_params["aspectRatio"] = image_config["aspectRatio"] + if "imageSize" in image_config: + mapped_params["imageSize"] = image_config["imageSize"] if ( "imageConfig" in supported_params diff --git a/litellm/types/llms/gemini.py b/litellm/types/llms/gemini.py index 9e3fea1bbbbe..b1aba2e70f66 100644 --- a/litellm/types/llms/gemini.py +++ b/litellm/types/llms/gemini.py @@ -171,6 +171,9 @@ class GeminiImageGenerationParameters(BaseModel): aspectRatio: Optional[str] = None """Aspect ratio for generated images (e.g., '1:1', '16:9', '9:16', '4:3', '3:4')""" + imageSize: Optional[str] = None + """Image size for generated images (e.g., '512', '1K', '2K', '4K')""" + personGeneration: Optional[str] = None """Controls person generation in images""" diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index b3463863ba7c..0a5aebdf91b9 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -70,7 +70,7 @@ ("1376x768", "16:9", "1K"), ("2752x1536", "16:9", "2K"), ("5504x3072", "16:9", "4K"), - ("792x168", "21:9", "512"), + ("792x336", "21:9", "512"), ("1584x672", "21:9", "1K"), ("3168x1344", "21:9", "2K"), ("6336x2688", "21:9", "4K"), @@ -609,6 +609,7 @@ def test_gemini_imagen_models_use_predict_endpoint(): assert "instances" in request_data assert "parameters" in request_data assert request_data["parameters"]["aspectRatio"] == "4:3" + assert request_data["parameters"]["imageSize"] == "1K" assert "imageConfig" not in request_data["parameters"] diff --git a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py index 71afcc9a3b45..73aa4e3b343b 100644 --- a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py +++ b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py @@ -7,6 +7,7 @@ import httpx import pytest +import litellm from litellm.llms.gemini.image_edit.transformation import GeminiImageEditConfig @@ -74,6 +75,22 @@ def test_map_openai_params_parses_form_image_config_json(self) -> None: assert mapped["imageConfig"] == {"aspectRatio": "16:9", "imageSize": "1K"} + def test_map_openai_params_rejects_malformed_form_image_config_json( + self, + ) -> None: + optional_params: Dict[str, object] = { + "imageConfig": "{bad", + } + + with pytest.raises(litellm.UnsupportedParamsError) as exc_info: + self.config.map_openai_params( + image_edit_optional_params=optional_params, # type: ignore[arg-type] + model="gemini-3-pro-image-preview", + drop_params=False, + ) + + assert "`imageConfig` must be valid JSON" in str(exc_info.value) + def test_transform_image_edit_request(self) -> None: image_bytes = b"fake_image_data" image = BytesIO(image_bytes) From 5902dffebc08bf3e1a0d8ed3e2464755d3859016 Mon Sep 17 00:00:00 2001 From: Dmitriy Alergant Date: Mon, 18 May 2026 18:38:50 -0400 Subject: [PATCH 03/12] Format Gemini image generation transform --- litellm/llms/gemini/image_generation/transformation.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index d012037bca3c..0f9b788f377f 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -70,9 +70,8 @@ def map_openai_params( if "imageSize" in image_config: mapped_params["imageSize"] = image_config["imageSize"] - if ( - "imageConfig" in supported_params - and isinstance(non_default_params.get("imageConfig"), dict) + if "imageConfig" in supported_params and isinstance( + non_default_params.get("imageConfig"), dict ): mapped_params["imageConfig"] = non_default_params["imageConfig"] From bed21171c9a358f934751adea92d0938de9ca82e Mon Sep 17 00:00:00 2001 From: Dmitriy Alergant Date: Tue, 19 May 2026 13:23:55 -0400 Subject: [PATCH 04/12] Fix Gemini image token usage logging --- .../llms/gemini/image_edit/cost_calculator.py | 23 ++----- .../llms/gemini/image_edit/transformation.py | 14 +++- .../gemini/image_generation/transformation.py | 42 ++--------- .../llms/gemini/image_usage_transformation.py | 40 +++++++++++ .../test_gemini_image_edit_transformation.py | 25 ++++++- .../llms/gemini/test_cost_calculator.py | 69 ++++++++++++++++++- ..._gemini_image_generation_transformation.py | 63 +++++++++++++++++ 7 files changed, 220 insertions(+), 56 deletions(-) create mode 100644 litellm/llms/gemini/image_usage_transformation.py create mode 100644 tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py diff --git a/litellm/llms/gemini/image_edit/cost_calculator.py b/litellm/llms/gemini/image_edit/cost_calculator.py index 2e332a7fc00d..956edb849a09 100644 --- a/litellm/llms/gemini/image_edit/cost_calculator.py +++ b/litellm/llms/gemini/image_edit/cost_calculator.py @@ -4,8 +4,9 @@ from typing import Any -import litellm -from litellm.types.utils import ImageResponse +from litellm.llms.gemini.image_generation.cost_calculator import ( + cost_calculator as image_generation_cost_calculator, +) def cost_calculator( @@ -15,20 +16,10 @@ def cost_calculator( """ Gemini image edit cost calculator. - Mirrors image generation pricing: charge per returned image based on - model metadata (`output_cost_per_image`). + Gemini image edits and generations share image response billing behavior: + use provider token usage when present, otherwise fall back to per-image pricing. """ - model_info = litellm.get_model_info( + return image_generation_cost_calculator( model=model, - custom_llm_provider="gemini", + image_response=image_response, ) - - output_cost_per_image: float = model_info.get("output_cost_per_image") or 0.0 - - if not isinstance(image_response, ImageResponse): - raise ValueError( - f"image_response must be of type ImageResponse got type={type(image_response)}" - ) - - num_images = len(image_response.data or []) - return output_cost_per_image * num_images diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index 75115b2653b2..fb3cbd201f2b 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -13,10 +13,18 @@ map_openai_size_to_gemini_image_config, supports_gemini_image_size, ) +from litellm.llms.gemini.image_usage_transformation import ( + transform_gemini_image_usage, +) from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import FileTypes, ImageObject, ImageResponse, OpenAIImage +from litellm.types.utils import ( + FileTypes, + ImageObject, + ImageResponse, + OpenAIImage, +) if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -176,6 +184,10 @@ def transform_image_edit_response( ) model_response.data = cast(List[OpenAIImage], data_list) + if "usageMetadata" in response_json: + model_response.usage = transform_gemini_image_usage( + response_json["usageMetadata"] + ) return model_response def _prepare_inline_image_parts( diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 0f9b788f377f..79be708a0919 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -9,18 +9,16 @@ map_openai_size_to_gemini_image_config, supports_gemini_image_size, ) +from litellm.llms.gemini.image_usage_transformation import ( + transform_gemini_image_usage, +) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.gemini import GeminiImageGenerationRequest from litellm.types.llms.openai import ( AllMessageValues, OpenAIImageGenerationOptionalParams, ) -from litellm.types.utils import ( - ImageObject, - ImageResponse, - ImageUsage, - ImageUsageInputTokensDetails, -) +from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -84,36 +82,6 @@ def map_openai_params( mapped_params[k] = v return mapped_params - def _transform_image_usage(self, usage_metadata: dict) -> ImageUsage: - """ - Transform Gemini usageMetadata to ImageUsage format - """ - input_tokens_details = ImageUsageInputTokensDetails( - image_tokens=0, - text_tokens=0, - ) - - # Extract detailed token counts from promptTokensDetails - tokens_details = usage_metadata.get("promptTokensDetails", []) - for details in tokens_details: - if isinstance(details, dict): - modality = str(details.get("modality", "")).upper() - raw_token_count = details.get( - "tokenCount", details.get("token_count", 0) - ) - token_count = raw_token_count if isinstance(raw_token_count, int) else 0 - if modality == "TEXT": - input_tokens_details.text_tokens += token_count - elif modality == "IMAGE": - input_tokens_details.image_tokens += token_count - - return ImageUsage( - input_tokens=usage_metadata.get("promptTokenCount", 0), - input_tokens_details=input_tokens_details, - output_tokens=usage_metadata.get("candidatesTokenCount", 0), - total_tokens=usage_metadata.get("totalTokenCount", 0), - ) - def get_complete_url( self, api_base: Optional[str], @@ -278,7 +246,7 @@ def transform_image_generation_response( # Extract usage metadata for Gemini models if "usageMetadata" in response_data: - model_response.usage = self._transform_image_usage( + model_response.usage = transform_gemini_image_usage( response_data["usageMetadata"] ) else: diff --git a/litellm/llms/gemini/image_usage_transformation.py b/litellm/llms/gemini/image_usage_transformation.py new file mode 100644 index 000000000000..882d0adc96dc --- /dev/null +++ b/litellm/llms/gemini/image_usage_transformation.py @@ -0,0 +1,40 @@ +from litellm.types.utils import ImageUsage, ImageUsageInputTokensDetails + + +def transform_gemini_image_usage(usage_metadata: dict) -> ImageUsage: + """ + Transform Gemini usageMetadata to ImageUsage format. + """ + input_tokens_details = ImageUsageInputTokensDetails( + image_tokens=0, + text_tokens=0, + ) + + for details in usage_metadata.get("promptTokensDetails", []): + if isinstance(details, dict): + modality = str(details.get("modality", "")).upper() + raw_token_count = details.get("tokenCount", details.get("token_count", 0)) + token_count = raw_token_count if isinstance(raw_token_count, int) else 0 + if modality == "TEXT": + input_tokens_details.text_tokens += token_count + elif modality == "IMAGE": + input_tokens_details.image_tokens += token_count + + output_tokens = usage_metadata.get("candidatesTokenCount", 0) + return ImageUsage( + input_tokens=usage_metadata.get("promptTokenCount", 0), + input_tokens_details=input_tokens_details, + output_tokens=output_tokens, + total_tokens=usage_metadata.get("totalTokenCount", 0), + prompt_tokens=usage_metadata.get("promptTokenCount", 0), + prompt_tokens_details=input_tokens_details.model_dump(), + completion_tokens=output_tokens, + completion_tokens_details={ + "text_tokens": 0, + "image_tokens": output_tokens, + }, + output_tokens_details={ + "text_tokens": 0, + "image_tokens": output_tokens, + }, + ) diff --git a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py index 73aa4e3b343b..fc76f5a821e5 100644 --- a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py +++ b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py @@ -8,6 +8,7 @@ import pytest import litellm +from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.gemini.image_edit.transformation import GeminiImageEditConfig @@ -192,7 +193,16 @@ def test_transform_image_edit_response(self) -> None: ] } }, - ] + ], + "usageMetadata": { + "promptTokenCount": 35, + "candidatesTokenCount": 1716, + "totalTokenCount": 1751, + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": 30}, + {"modality": "IMAGE", "tokenCount": 5}, + ], + }, } mock_response = MagicMock(spec=httpx.Response) @@ -215,6 +225,19 @@ def test_transform_image_edit_response(self) -> None: "utf-8" ) + usage = image_response.model_dump()["usage"] + assert usage["input_tokens"] == 35 + assert usage["output_tokens"] == 1716 + assert usage["prompt_tokens"] == 35 + assert usage["completion_tokens"] == 1716 + assert usage["prompt_tokens_details"]["image_tokens"] == 5 + assert usage["completion_tokens_details"]["image_tokens"] == 1716 + + logging_usage = StandardLoggingPayloadSetup.get_usage_as_dict( + response_obj=image_response.model_dump() + ) + assert logging_usage["completion_tokens_details"]["image_tokens"] == 1716 + def test_transform_image_edit_request_without_image_raises(self) -> None: optional_params = {} diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/test_litellm/llms/gemini/test_cost_calculator.py index 9bb83aa7cff0..e8a3fc25a411 100644 --- a/tests/test_litellm/llms/gemini/test_cost_calculator.py +++ b/tests/test_litellm/llms/gemini/test_cost_calculator.py @@ -1,7 +1,20 @@ +import os + import pytest +import litellm from litellm.llms.gemini.cost_calculator import cost_per_web_search_request -from litellm.types.utils import PromptTokensDetailsWrapper, Usage +from litellm.llms.gemini.image_edit.cost_calculator import ( + cost_calculator as gemini_image_edit_cost_calculator, +) +from litellm.types.utils import ( + ImageObject, + ImageResponse, + ImageUsage, + ImageUsageInputTokensDetails, + PromptTokensDetailsWrapper, + Usage, +) def _make_usage(web_search_requests: int) -> Usage: @@ -63,3 +76,57 @@ def test_no_usage_details(): usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) cost = cost_per_web_search_request(usage=usage, model_info=model_info) assert cost == 0.0 + + +def test_gemini_image_edit_cost_prefers_token_usage_metadata(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + model = "gemini/gemini-3-pro-image-preview" + model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") + + input_text_tokens = 20 + input_image_tokens = 1120 + output_image_tokens = 1120 + prompt_tokens = input_text_tokens + input_image_tokens + image_response = ImageResponse( + data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")], + usage=ImageUsage( + input_tokens=prompt_tokens, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=input_text_tokens, + image_tokens=input_image_tokens, + ), + output_tokens=output_image_tokens, + total_tokens=prompt_tokens + output_image_tokens, + ), + ) + + cost = gemini_image_edit_cost_calculator( + model=model, + image_response=image_response, + ) + + expected_cost = ( + prompt_tokens * model_info["input_cost_per_token"] + + output_image_tokens * model_info["output_cost_per_image_token"] + ) + flat_image_cost = len(image_response.data or []) * model_info["output_cost_per_image"] + assert round(cost, 10) == round(expected_cost, 10) + assert cost != flat_image_cost + + +def test_gemini_image_edit_cost_falls_back_to_flat_image_pricing(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + model = "gemini/gemini-3-pro-image-preview" + model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") + image_response = ImageResponse( + data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")] + ) + + cost = gemini_image_edit_cost_calculator( + model=model, + image_response=image_response, + ) + + assert cost == len(image_response.data or []) * model_info["output_cost_per_image"] diff --git a/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py b/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py new file mode 100644 index 000000000000..478013a4b038 --- /dev/null +++ b/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py @@ -0,0 +1,63 @@ +import httpx + +from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup +from litellm.llms.gemini.image_generation.transformation import GoogleImageGenConfig +from litellm.types.utils import ImageResponse + + +def test_gemini_image_generation_usage_includes_chat_token_details(): + config = GoogleImageGenConfig() + raw_response = httpx.Response( + status_code=200, + json={ + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "fake-image", + } + } + ] + } + } + ], + "usageMetadata": { + "promptTokenCount": 35, + "candidatesTokenCount": 1716, + "totalTokenCount": 1751, + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": 30}, + {"modality": "IMAGE", "tokenCount": 5}, + ], + }, + }, + ) + + result = config.transform_image_generation_response( + model="gemini-3.1-flash-image-preview", + raw_response=raw_response, + model_response=ImageResponse(data=[]), + logging_obj=None, + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + usage = result.model_dump()["usage"] + + assert usage["input_tokens"] == 35 + assert usage["output_tokens"] == 1716 + assert usage["prompt_tokens"] == 35 + assert usage["completion_tokens"] == 1716 + assert usage["prompt_tokens_details"]["image_tokens"] == 5 + assert usage["completion_tokens_details"]["image_tokens"] == 1716 + assert usage["output_tokens_details"]["image_tokens"] == 1716 + + logging_usage = StandardLoggingPayloadSetup.get_usage_as_dict( + response_obj=result.model_dump() + ) + assert logging_usage["completion_tokens_details"]["image_tokens"] == 1716 From 8dcf88c5ed643192a377332755c81259b667addc Mon Sep 17 00:00:00 2001 From: Dmitriy Alergant Date: Tue, 19 May 2026 13:34:40 -0400 Subject: [PATCH 05/12] Share Gemini image request helpers --- litellm/llms/gemini/common_utils.py | 78 +++++++++++++++++++ .../llms/gemini/image_edit/transformation.py | 59 ++++---------- .../gemini/image_generation/transformation.py | 60 +++----------- .../test_gemini_image_edit_transformation.py | 3 + ..._gemini_image_generation_transformation.py | 42 ++++++++++ 5 files changed, 148 insertions(+), 94 deletions(-) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 7c7c38d760d3..cb65a0f0f9f6 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -1,5 +1,6 @@ import base64 import datetime +import json import math from typing import Any, Dict, List, Optional, Union @@ -115,6 +116,83 @@ def supports_gemini_image_size(model: str) -> bool: return "2.5-flash" not in model +def map_openai_image_params_to_gemini( + params: Dict[str, Any], + model: str, + supported_params: List[str], + optional_params: Optional[Dict[str, Any]] = None, + parse_image_config_string: bool = False, +) -> Dict[str, Any]: + optional_params = optional_params or {} + filtered_params = { + key: value for key, value in params.items() if key in supported_params + } + + mapped_params: Dict[str, Any] = {} + + if "n" in filtered_params and "n" not in optional_params: + mapped_params["sampleCount"] = filtered_params["n"] + + if "size" in filtered_params and "size" not in optional_params: + image_config = map_openai_size_to_gemini_image_config( + filtered_params["size"], + model, + ) + if image_config is not None: + if "gemini" in model: + mapped_params["imageConfig"] = image_config + else: + mapped_params["aspectRatio"] = image_config["aspectRatio"] + if "imageSize" in image_config: + mapped_params["imageSize"] = image_config["imageSize"] + + image_config_param = filtered_params.get("imageConfig") + if isinstance(image_config_param, str) and parse_image_config_string: + try: + image_config_param = json.loads(image_config_param) + except json.JSONDecodeError as exc: + raise litellm.UnsupportedParamsError( + model=model, + message="`imageConfig` must be valid JSON when provided as a string.", + ) from exc + if isinstance(image_config_param, dict): + mapped_params["imageConfig"] = image_config_param + + for key, value in filtered_params.items(): + if key not in ("n", "size", "imageConfig") and key not in optional_params: + mapped_params[key] = value + + return mapped_params + + +def get_gemini_image_generation_config( + model: str, + optional_params: Dict[str, Any], +) -> Dict[str, Any]: + generation_config: Dict[str, Any] = {"response_modalities": ["IMAGE", "TEXT"]} + + image_config: Dict[str, Any] = {} + if isinstance(optional_params.get("imageConfig"), dict): + image_config.update(optional_params["imageConfig"]) + + if not supports_gemini_image_size(model): + image_config.pop("imageSize", None) + + if image_config: + generation_config["imageConfig"] = image_config + + candidate_count = ( + optional_params.get("candidateCount") + or optional_params.get("candidate_count") + or optional_params.get("sampleCount") + or optional_params.get("n") + ) + if candidate_count is not None: + generation_config["candidateCount"] = candidate_count + + return generation_config + + def _parse_openai_image_size(size: str) -> Optional[tuple[int, int]]: if size == "auto": return None diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index fb3cbd201f2b..2316361d6e7f 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -1,17 +1,15 @@ import base64 -import json from io import BufferedReader, BytesIO from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast import httpx from httpx._types import RequestFiles -import litellm from litellm.images.utils import ImageEditRequestUtils from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.gemini.common_utils import ( - map_openai_size_to_gemini_image_config, - supports_gemini_image_size, + get_gemini_image_generation_config, + map_openai_image_params_to_gemini, ) from litellm.llms.gemini.image_usage_transformation import ( transform_gemini_image_usage, @@ -36,7 +34,7 @@ class GeminiImageEditConfig(BaseImageEditConfig): DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta" - SUPPORTED_PARAMS: List[str] = ["size", "imageConfig"] + SUPPORTED_PARAMS: List[str] = ["n", "size", "imageConfig"] def get_supported_openai_params(self, model: str) -> List[str]: return list(self.SUPPORTED_PARAMS) @@ -47,36 +45,12 @@ def map_openai_params( model: str, drop_params: bool, ) -> Dict[str, Any]: - supported_params = self.get_supported_openai_params(model) - filtered_params = { - key: value - for key, value in image_edit_optional_params.items() - if key in supported_params - } - - mapped_params: Dict[str, Any] = {} - - if "size" in filtered_params: - image_config = map_openai_size_to_gemini_image_config( - filtered_params["size"], # type: ignore[arg-type] - model, - ) - if image_config is not None: - mapped_params["imageConfig"] = image_config - - image_config_param = filtered_params.get("imageConfig") - if isinstance(image_config_param, str): - try: - image_config_param = json.loads(image_config_param) - except json.JSONDecodeError as exc: - raise litellm.UnsupportedParamsError( - model=model, - message="`imageConfig` must be valid JSON when provided as a string.", - ) from exc - if isinstance(image_config_param, dict): - mapped_params["imageConfig"] = image_config_param - - return mapped_params + return map_openai_image_params_to_gemini( + params=image_edit_optional_params, # type: ignore[arg-type] + model=model, + supported_params=self.get_supported_openai_params(model), + parse_image_config_string=True, + ) def validate_environment( self, @@ -136,17 +110,10 @@ def transform_image_edit_request( # type: ignore[override] request_body: Dict[str, Any] = {"contents": contents} - generation_config: Dict[str, Any] = {} - - if isinstance(image_edit_optional_request_params.get("imageConfig"), dict): - image_config = dict(image_edit_optional_request_params["imageConfig"]) - if not supports_gemini_image_size(model): - image_config.pop("imageSize", None) - if image_config: - generation_config["imageConfig"] = image_config - - if generation_config: - request_body["generationConfig"] = generation_config + request_body["generationConfig"] = get_gemini_image_generation_config( + model=model, + optional_params=image_edit_optional_request_params, + ) empty_files = cast(RequestFiles, []) return request_body, empty_files diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 79be708a0919..95d80b4d11d2 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -6,8 +6,8 @@ BaseImageGenerationConfig, ) from litellm.llms.gemini.common_utils import ( - map_openai_size_to_gemini_image_config, - supports_gemini_image_size, + get_gemini_image_generation_config, + map_openai_image_params_to_gemini, ) from litellm.llms.gemini.image_usage_transformation import ( transform_gemini_image_usage, @@ -50,37 +50,12 @@ def map_openai_params( model: str, drop_params: bool, ) -> dict: - supported_params = self.get_supported_openai_params(model) - mapped_params = {} - - if "n" in non_default_params and "n" not in optional_params: - mapped_params["sampleCount"] = non_default_params["n"] - - if "size" in non_default_params and "size" not in optional_params: - image_config = map_openai_size_to_gemini_image_config( - non_default_params["size"], model - ) - if image_config is not None: - if "gemini" in model: - mapped_params["imageConfig"] = image_config - else: - mapped_params["aspectRatio"] = image_config["aspectRatio"] - if "imageSize" in image_config: - mapped_params["imageSize"] = image_config["imageSize"] - - if "imageConfig" in supported_params and isinstance( - non_default_params.get("imageConfig"), dict - ): - mapped_params["imageConfig"] = non_default_params["imageConfig"] - - for k, v in non_default_params.items(): - if ( - k not in ("n", "size", "imageConfig") - and k not in optional_params - and k in supported_params - ): - mapped_params[k] = v - return mapped_params + return map_openai_image_params_to_gemini( + params=non_default_params, + model=model, + supported_params=self.get_supported_openai_params(model), + optional_params=optional_params, + ) def get_complete_url( self, @@ -157,23 +132,12 @@ def transform_image_generation_request( """ # For Gemini Flash Image Preview models, use standard Gemini format if "gemini" in model: - generation_config: Dict[str, Any] = { - "response_modalities": ["IMAGE", "TEXT"] - } - image_config: Dict[str, Any] = {} - - if isinstance(optional_params.get("imageConfig"), dict): - image_config.update(optional_params["imageConfig"]) - - if not supports_gemini_image_size(model): - image_config.pop("imageSize", None) - - if image_config: - generation_config["imageConfig"] = image_config - request_body: dict = { "contents": [{"parts": [{"text": prompt}]}], - "generationConfig": generation_config, + "generationConfig": get_gemini_image_generation_config( + model=model, + optional_params=optional_params, + ), } return request_body else: diff --git a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py index fc76f5a821e5..9b57e1991def 100644 --- a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py +++ b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py @@ -21,6 +21,7 @@ def setup_method(self) -> None: def test_map_openai_params(self) -> None: optional_params: Dict[str, object] = { + "n": 2, "size": "1792x1024", "response_format": "b64_json", "quality": "high", @@ -33,6 +34,7 @@ def test_map_openai_params(self) -> None: ) assert mapped["imageConfig"] == {"aspectRatio": "16:9"} + assert mapped["sampleCount"] == 2 assert "response_format" not in mapped assert "quality" not in mapped @@ -119,6 +121,7 @@ def test_transform_image_edit_request(self) -> None: assert base64.b64decode(inline_data["data"]) == image_bytes generation_config = request_body["generationConfig"] + assert generation_config["candidateCount"] == 2 assert generation_config["imageConfig"]["aspectRatio"] == "16:9" assert generation_config["imageConfig"]["imageSize"] == "2K" diff --git a/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py b/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py index 478013a4b038..f792c6406760 100644 --- a/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py +++ b/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py @@ -5,6 +5,48 @@ from litellm.types.utils import ImageResponse +def test_gemini_image_generation_request_uses_shared_generation_config(): + config = GoogleImageGenConfig() + + request = config.transform_image_generation_request( + model="gemini-3.1-flash-image-preview", + prompt="Generate a simple app icon", + optional_params={ + "sampleCount": 2, + "imageConfig": {"aspectRatio": "16:9", "imageSize": "2K"}, + }, + litellm_params={}, + headers={}, + ) + + assert request["contents"][0]["parts"] == [{"text": "Generate a simple app icon"}] + assert request["generationConfig"] == { + "response_modalities": ["IMAGE", "TEXT"], + "imageConfig": {"aspectRatio": "16:9", "imageSize": "2K"}, + "candidateCount": 2, + } + + +def test_gemini_image_generation_map_openai_params_maps_n_size_and_image_config(): + config = GoogleImageGenConfig() + + mapped = config.map_openai_params( + non_default_params={ + "n": 2, + "size": "768x1376", + "imageConfig": {"aspectRatio": "1:1", "imageSize": "512"}, + }, + optional_params={}, + model="gemini-3.1-flash-image-preview", + drop_params=False, + ) + + assert mapped == { + "sampleCount": 2, + "imageConfig": {"aspectRatio": "1:1", "imageSize": "512"}, + } + + def test_gemini_image_generation_usage_includes_chat_token_details(): config = GoogleImageGenConfig() raw_response = httpx.Response( From 75f887f372cee91f3a795fa5883465dae4b6d980 Mon Sep 17 00:00:00 2001 From: Dmitriy Alergant Date: Tue, 19 May 2026 13:43:31 -0400 Subject: [PATCH 06/12] Fix Gemini Imagen model routing --- litellm/llms/gemini/common_utils.py | 7 ++- .../gemini/image_generation/transformation.py | 9 +-- ..._gemini_image_generation_transformation.py | 57 +++++++++++++++++++ 3 files changed, 68 insertions(+), 5 deletions(-) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index cb65a0f0f9f6..6dcd5318de7b 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -116,6 +116,11 @@ def supports_gemini_image_size(model: str) -> bool: return "2.5-flash" not in model +def is_gemini_image_model(model: str) -> bool: + base_model = model.split("/", 1)[-1] + return "gemini" in base_model + + def map_openai_image_params_to_gemini( params: Dict[str, Any], model: str, @@ -139,7 +144,7 @@ def map_openai_image_params_to_gemini( model, ) if image_config is not None: - if "gemini" in model: + if is_gemini_image_model(model): mapped_params["imageConfig"] = image_config else: mapped_params["aspectRatio"] = image_config["aspectRatio"] diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 95d80b4d11d2..5c8772c57383 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -7,6 +7,7 @@ ) from litellm.llms.gemini.common_utils import ( get_gemini_image_generation_config, + is_gemini_image_model, map_openai_image_params_to_gemini, ) from litellm.llms.gemini.image_usage_transformation import ( @@ -39,7 +40,7 @@ def get_supported_openai_params( https://ai.google.dev/gemini-api/docs/imagen """ supported_params = ["n", "size"] - if "gemini" in model: + if is_gemini_image_model(model): supported_params.append("imageConfig") return supported_params # type: ignore[return-value] @@ -79,7 +80,7 @@ def get_complete_url( complete_url = complete_url.rstrip("/") # Gemini Flash Image Preview models use generateContent endpoint - if "gemini" in model: + if is_gemini_image_model(model): complete_url = f"{complete_url}/models/{model}:generateContent" else: # All other Imagen models use predict endpoint @@ -131,7 +132,7 @@ def transform_image_generation_request( } """ # For Gemini Flash Image Preview models, use standard Gemini format - if "gemini" in model: + if is_gemini_image_model(model): request_body: dict = { "contents": [{"parts": [{"text": prompt}]}], "generationConfig": get_gemini_image_generation_config( @@ -184,7 +185,7 @@ def transform_image_generation_response( model_response.data = [] # Handle different response formats based on model - if "gemini" in model: + if is_gemini_image_model(model): # Gemini Flash Image Preview models return in candidates format candidates = response_data.get("candidates", []) for candidate in candidates: diff --git a/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py b/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py index f792c6406760..541350c75a40 100644 --- a/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py +++ b/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py @@ -47,6 +47,63 @@ def test_gemini_image_generation_map_openai_params_maps_n_size_and_image_config( } +def test_imagen_generation_with_provider_prefix_uses_imagen_params_and_response(): + config = GoogleImageGenConfig() + + mapped = config.map_openai_params( + non_default_params={ + "n": 1, + "size": "1024x1024", + }, + optional_params={}, + model="gemini/imagen-4.0-generate-001", + drop_params=False, + ) + assert mapped == { + "sampleCount": 1, + "aspectRatio": "1:1", + "imageSize": "1K", + } + + request = config.transform_image_generation_request( + model="gemini/imagen-4.0-generate-001", + prompt="Generate a simple app icon", + optional_params=mapped, + litellm_params={}, + headers={}, + ) + assert request == { + "instances": [{"prompt": "Generate a simple app icon"}], + "parameters": { + "sampleCount": 1, + "aspectRatio": "1:1", + "imageSize": "1K", + }, + } + + result = config.transform_image_generation_response( + model="gemini/imagen-4.0-generate-001", + raw_response=httpx.Response( + status_code=200, + json={ + "predictions": [ + { + "bytesBase64Encoded": "fake-imagen-image", + } + ] + }, + ), + model_response=ImageResponse(data=[]), + logging_obj=None, + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert result.data is not None + assert result.data[0].b64_json == "fake-imagen-image" + + def test_gemini_image_generation_usage_includes_chat_token_details(): config = GoogleImageGenConfig() raw_response = httpx.Response( From 4bd8a35a583324e420807477d4ad4be240720846 Mon Sep 17 00:00:00 2001 From: Dmitriy Alergant Date: Tue, 19 May 2026 14:41:34 -0400 Subject: [PATCH 07/12] Fixes as per self code review --- .../litellm_core_utils/llm_cost_calc/utils.py | 50 +++++++++-- litellm/llms/gemini/common_utils.py | 10 ++- .../llms/gemini/image_usage_transformation.py | 56 +++++++++---- litellm/types/llms/gemini.py | 6 +- .../llms/gemini/test_cost_calculator.py | 56 +++++++++++++ ..._gemini_image_generation_transformation.py | 84 ++++++++++++++++++- 6 files changed, 230 insertions(+), 32 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 59d0465e6d43..009d093494da 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -30,6 +30,14 @@ ) +def _get_token_detail_value(details: object, key: str) -> Optional[int]: + if isinstance(details, dict): + value = details.get(key) + else: + value = getattr(details, key, None) + return value if isinstance(value, int) else None + + def _is_above_128k(tokens: float) -> bool: if tokens > 128000: return True @@ -821,17 +829,47 @@ def calculate_image_response_cost_from_usage( cached_tokens=0, ) + output_tokens_details = getattr(usage, "completion_tokens_details", None) + if output_tokens_details is None: + output_tokens_details = getattr(usage, "output_tokens_details", None) + + if output_tokens_details is None: + completion_tokens_details = CompletionTokensDetailsWrapper( + text_tokens=0, + image_tokens=completion_tokens, + reasoning_tokens=0, + audio_tokens=0, + ) + else: + text_tokens = _get_token_detail_value(output_tokens_details, "text_tokens") or 0 + image_tokens = ( + _get_token_detail_value(output_tokens_details, "image_tokens") or 0 + ) + audio_tokens = ( + _get_token_detail_value(output_tokens_details, "audio_tokens") or 0 + ) + reasoning_tokens = ( + _get_token_detail_value(output_tokens_details, "reasoning_tokens") or 0 + ) + known_output_tokens = ( + text_tokens + image_tokens + audio_tokens + reasoning_tokens + ) + if completion_tokens > known_output_tokens: + text_tokens += completion_tokens - known_output_tokens + + completion_tokens_details = CompletionTokensDetailsWrapper( + text_tokens=text_tokens, + image_tokens=image_tokens, + reasoning_tokens=reasoning_tokens, + audio_tokens=audio_tokens, + ) + normalized_usage = Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens, prompt_tokens_details=prompt_tokens_details, - completion_tokens_details=CompletionTokensDetailsWrapper( - text_tokens=0, - image_tokens=completion_tokens, - reasoning_tokens=0, - audio_tokens=0, - ), + completion_tokens_details=completion_tokens_details, ) prompt_cost, completion_cost = generic_cost_per_token( diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 6dcd5318de7b..805493e9737a 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -93,6 +93,8 @@ (896, 1280): "3:4", } +IMAGEN_SUPPORTED_IMAGE_SIZES = {"1K", "2K"} + def map_openai_size_to_gemini_image_config( size: str, model: str @@ -105,8 +107,12 @@ def map_openai_size_to_gemini_image_config( image_config = { "aspectRatio": _map_dimensions_to_gemini_aspect_ratio(width, height) } - if supports_gemini_image_size(model): - image_config["imageSize"] = _map_dimensions_to_gemini_image_size(width, height) + image_size = _map_dimensions_to_gemini_image_size(width, height) + if is_gemini_image_model(model): + if supports_gemini_image_size(model): + image_config["imageSize"] = image_size + elif image_size in IMAGEN_SUPPORTED_IMAGE_SIZES: + image_config["imageSize"] = image_size return image_config diff --git a/litellm/llms/gemini/image_usage_transformation.py b/litellm/llms/gemini/image_usage_transformation.py index 882d0adc96dc..0c500186b156 100644 --- a/litellm/llms/gemini/image_usage_transformation.py +++ b/litellm/llms/gemini/image_usage_transformation.py @@ -1,26 +1,52 @@ from litellm.types.utils import ImageUsage, ImageUsageInputTokensDetails -def transform_gemini_image_usage(usage_metadata: dict) -> ImageUsage: - """ - Transform Gemini usageMetadata to ImageUsage format. - """ - input_tokens_details = ImageUsageInputTokensDetails( +def _get_token_count(details: dict) -> int: + raw_token_count = details.get("tokenCount", details.get("token_count", 0)) + return raw_token_count if isinstance(raw_token_count, int) else 0 + + +def _sum_modality_token_details( + usage_metadata: dict, details_key: str +) -> ImageUsageInputTokensDetails: + tokens_details = ImageUsageInputTokensDetails( image_tokens=0, text_tokens=0, ) - for details in usage_metadata.get("promptTokensDetails", []): + for details in usage_metadata.get(details_key, []): if isinstance(details, dict): modality = str(details.get("modality", "")).upper() - raw_token_count = details.get("tokenCount", details.get("token_count", 0)) - token_count = raw_token_count if isinstance(raw_token_count, int) else 0 + token_count = _get_token_count(details) if modality == "TEXT": - input_tokens_details.text_tokens += token_count + tokens_details.text_tokens += token_count elif modality == "IMAGE": - input_tokens_details.image_tokens += token_count + tokens_details.image_tokens += token_count + + return tokens_details + +def transform_gemini_image_usage(usage_metadata: dict) -> ImageUsage: + """ + Transform Gemini usageMetadata to ImageUsage format. + """ + input_tokens_details = _sum_modality_token_details( + usage_metadata, "promptTokensDetails" + ) output_tokens = usage_metadata.get("candidatesTokenCount", 0) + output_tokens_details = _sum_modality_token_details( + usage_metadata, "candidatesTokensDetails" + ) + + if not usage_metadata.get("candidatesTokensDetails"): + output_tokens_details.image_tokens = output_tokens + else: + known_output_tokens = ( + output_tokens_details.text_tokens + output_tokens_details.image_tokens + ) + if output_tokens > known_output_tokens: + output_tokens_details.text_tokens += output_tokens - known_output_tokens + return ImageUsage( input_tokens=usage_metadata.get("promptTokenCount", 0), input_tokens_details=input_tokens_details, @@ -29,12 +55,6 @@ def transform_gemini_image_usage(usage_metadata: dict) -> ImageUsage: prompt_tokens=usage_metadata.get("promptTokenCount", 0), prompt_tokens_details=input_tokens_details.model_dump(), completion_tokens=output_tokens, - completion_tokens_details={ - "text_tokens": 0, - "image_tokens": output_tokens, - }, - output_tokens_details={ - "text_tokens": 0, - "image_tokens": output_tokens, - }, + completion_tokens_details=output_tokens_details.model_dump(), + output_tokens_details=output_tokens_details.model_dump(), ) diff --git a/litellm/types/llms/gemini.py b/litellm/types/llms/gemini.py index b1aba2e70f66..2ef4a9dc0b36 100644 --- a/litellm/types/llms/gemini.py +++ b/litellm/types/llms/gemini.py @@ -1,7 +1,7 @@ from enum import Enum -from typing import Any, Dict, Iterable, List, Literal, Optional, Union +from typing import Any, Dict, List, Literal, Optional -from typing_extensions import Required, TypedDict +from typing_extensions import TypedDict from .vertex_ai import ( GenerationConfig, @@ -172,7 +172,7 @@ class GeminiImageGenerationParameters(BaseModel): """Aspect ratio for generated images (e.g., '1:1', '16:9', '9:16', '4:3', '3:4')""" imageSize: Optional[str] = None - """Image size for generated images (e.g., '512', '1K', '2K', '4K')""" + """Image size for generated images (e.g., '1K', '2K')""" personGeneration: Optional[str] = None """Controls person generation in images""" diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/test_litellm/llms/gemini/test_cost_calculator.py index e8a3fc25a411..c0616b561d65 100644 --- a/tests/test_litellm/llms/gemini/test_cost_calculator.py +++ b/tests/test_litellm/llms/gemini/test_cost_calculator.py @@ -115,6 +115,62 @@ def test_gemini_image_edit_cost_prefers_token_usage_metadata(): assert cost != flat_image_cost +def test_gemini_image_edit_cost_uses_output_token_details(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + model = "gemini/gemini-3-pro-image-preview" + model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") + + input_text_tokens = 20 + output_text_tokens = 213 + output_image_tokens = 1120 + output_tokens = output_text_tokens + output_image_tokens + image_response = ImageResponse( + data=[ImageObject(b64_json="img1")], + usage=ImageUsage( + input_tokens=input_text_tokens, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=input_text_tokens, + image_tokens=0, + ), + output_tokens=output_tokens, + total_tokens=input_text_tokens + output_tokens, + prompt_tokens=input_text_tokens, + completion_tokens=output_tokens, + prompt_tokens_details={ + "text_tokens": input_text_tokens, + "image_tokens": 0, + }, + completion_tokens_details={ + "text_tokens": output_text_tokens, + "image_tokens": output_image_tokens, + }, + output_tokens_details={ + "text_tokens": output_text_tokens, + "image_tokens": output_image_tokens, + }, + ), + ) + + cost = gemini_image_edit_cost_calculator( + model=model, + image_response=image_response, + ) + + expected_cost = ( + input_text_tokens * model_info["input_cost_per_token"] + + output_text_tokens * model_info["output_cost_per_token"] + + output_image_tokens * model_info["output_cost_per_image_token"] + ) + all_output_as_image_cost = ( + input_text_tokens * model_info["input_cost_per_token"] + + (output_text_tokens + output_image_tokens) + * model_info["output_cost_per_image_token"] + ) + assert round(cost, 10) == round(expected_cost, 10) + assert cost != all_output_as_image_cost + + def test_gemini_image_edit_cost_falls_back_to_flat_image_pricing(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") diff --git a/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py b/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py index 541350c75a40..d2429cfc4ac1 100644 --- a/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py +++ b/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py @@ -104,6 +104,33 @@ def test_imagen_generation_with_provider_prefix_uses_imagen_params_and_response( assert result.data[0].b64_json == "fake-imagen-image" +def test_imagen_generation_omits_unsupported_openai_size_image_size(): + config = GoogleImageGenConfig() + + mapped = config.map_openai_params( + non_default_params={ + "size": "512x512", + }, + optional_params={}, + model="gemini/imagen-4.0-generate-001", + drop_params=False, + ) + assert mapped == {"aspectRatio": "1:1"} + + request = config.transform_image_generation_request( + model="gemini/imagen-4.0-generate-001", + prompt="Generate a simple app icon", + optional_params=mapped, + litellm_params={}, + headers={}, + ) + + assert request == { + "instances": [{"prompt": "Generate a simple app icon"}], + "parameters": {"aspectRatio": "1:1"}, + } + + def test_gemini_image_generation_usage_includes_chat_token_details(): config = GoogleImageGenConfig() raw_response = httpx.Response( @@ -131,6 +158,10 @@ def test_gemini_image_generation_usage_includes_chat_token_details(): {"modality": "TEXT", "tokenCount": 30}, {"modality": "IMAGE", "tokenCount": 5}, ], + "candidatesTokensDetails": [ + {"modality": "TEXT", "tokenCount": 213}, + {"modality": "IMAGE", "tokenCount": 1120}, + ], }, }, ) @@ -153,10 +184,57 @@ def test_gemini_image_generation_usage_includes_chat_token_details(): assert usage["prompt_tokens"] == 35 assert usage["completion_tokens"] == 1716 assert usage["prompt_tokens_details"]["image_tokens"] == 5 - assert usage["completion_tokens_details"]["image_tokens"] == 1716 - assert usage["output_tokens_details"]["image_tokens"] == 1716 + assert usage["completion_tokens_details"]["text_tokens"] == 596 + assert usage["completion_tokens_details"]["image_tokens"] == 1120 + assert usage["output_tokens_details"]["text_tokens"] == 596 + assert usage["output_tokens_details"]["image_tokens"] == 1120 logging_usage = StandardLoggingPayloadSetup.get_usage_as_dict( response_obj=result.model_dump() ) - assert logging_usage["completion_tokens_details"]["image_tokens"] == 1716 + assert logging_usage["completion_tokens_details"]["text_tokens"] == 596 + assert logging_usage["completion_tokens_details"]["image_tokens"] == 1120 + + +def test_gemini_image_generation_usage_without_output_details_treats_output_as_image(): + config = GoogleImageGenConfig() + raw_response = httpx.Response( + status_code=200, + json={ + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "fake-image", + } + } + ] + } + } + ], + "usageMetadata": { + "promptTokenCount": 35, + "candidatesTokenCount": 1716, + "totalTokenCount": 1751, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 35}], + }, + }, + ) + + result = config.transform_image_generation_response( + model="gemini-3.1-flash-image-preview", + raw_response=raw_response, + model_response=ImageResponse(data=[]), + logging_obj=None, + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + usage = result.model_dump()["usage"] + assert usage["completion_tokens_details"]["text_tokens"] == 0 + assert usage["completion_tokens_details"]["image_tokens"] == 1716 From ef3875e40a4d3d677079216a2c785f93f773e9d7 Mon Sep 17 00:00:00 2001 From: Dmitriy Alergant Date: Tue, 19 May 2026 15:15:12 -0400 Subject: [PATCH 08/12] Fixes per internal code review --- litellm/llms/gemini/common_utils.py | 7 ++- .../llms/gemini/image_usage_transformation.py | 20 ++++-- .../llms/gemini/test_cost_calculator.py | 63 ++++++++++++++++++- 3 files changed, 83 insertions(+), 7 deletions(-) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 805493e9737a..5ce223ec95df 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -111,7 +111,7 @@ def map_openai_size_to_gemini_image_config( if is_gemini_image_model(model): if supports_gemini_image_size(model): image_config["imageSize"] = image_size - elif image_size in IMAGEN_SUPPORTED_IMAGE_SIZES: + elif is_imagen_model(model) and image_size in IMAGEN_SUPPORTED_IMAGE_SIZES: image_config["imageSize"] = image_size return image_config @@ -127,6 +127,11 @@ def is_gemini_image_model(model: str) -> bool: return "gemini" in base_model +def is_imagen_model(model: str) -> bool: + base_model = model.split("/", 1)[-1] + return "imagen" in base_model + + def map_openai_image_params_to_gemini( params: Dict[str, Any], model: str, diff --git a/litellm/llms/gemini/image_usage_transformation.py b/litellm/llms/gemini/image_usage_transformation.py index 0c500186b156..d19e7a3f2e86 100644 --- a/litellm/llms/gemini/image_usage_transformation.py +++ b/litellm/llms/gemini/image_usage_transformation.py @@ -6,15 +6,23 @@ def _get_token_count(details: dict) -> int: return raw_token_count if isinstance(raw_token_count, int) else 0 +def _get_modality_token_details(usage_metadata: dict, *details_keys: str) -> list: + for details_key in details_keys: + details = usage_metadata.get(details_key) + if isinstance(details, list): + return details + return [] + + def _sum_modality_token_details( - usage_metadata: dict, details_key: str + usage_metadata: dict, *details_keys: str ) -> ImageUsageInputTokensDetails: tokens_details = ImageUsageInputTokensDetails( image_tokens=0, text_tokens=0, ) - for details in usage_metadata.get(details_key, []): + for details in _get_modality_token_details(usage_metadata, *details_keys): if isinstance(details, dict): modality = str(details.get("modality", "")).upper() token_count = _get_token_count(details) @@ -31,14 +39,16 @@ def transform_gemini_image_usage(usage_metadata: dict) -> ImageUsage: Transform Gemini usageMetadata to ImageUsage format. """ input_tokens_details = _sum_modality_token_details( - usage_metadata, "promptTokensDetails" + usage_metadata, "promptTokensDetails", "prompt_tokens_details" ) output_tokens = usage_metadata.get("candidatesTokenCount", 0) output_tokens_details = _sum_modality_token_details( - usage_metadata, "candidatesTokensDetails" + usage_metadata, "candidatesTokensDetails", "candidates_tokens_details" ) - if not usage_metadata.get("candidatesTokensDetails"): + if not _get_modality_token_details( + usage_metadata, "candidatesTokensDetails", "candidates_tokens_details" + ): output_tokens_details.image_tokens = output_tokens else: known_output_tokens = ( diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/test_litellm/llms/gemini/test_cost_calculator.py index c0616b561d65..6d51bcd2c889 100644 --- a/tests/test_litellm/llms/gemini/test_cost_calculator.py +++ b/tests/test_litellm/llms/gemini/test_cost_calculator.py @@ -7,6 +7,9 @@ from litellm.llms.gemini.image_edit.cost_calculator import ( cost_calculator as gemini_image_edit_cost_calculator, ) +from litellm.llms.gemini.image_generation.cost_calculator import ( + cost_calculator as gemini_image_generation_cost_calculator, +) from litellm.types.utils import ( ImageObject, ImageResponse, @@ -110,7 +113,9 @@ def test_gemini_image_edit_cost_prefers_token_usage_metadata(): prompt_tokens * model_info["input_cost_per_token"] + output_image_tokens * model_info["output_cost_per_image_token"] ) - flat_image_cost = len(image_response.data or []) * model_info["output_cost_per_image"] + flat_image_cost = ( + len(image_response.data or []) * model_info["output_cost_per_image"] + ) assert round(cost, 10) == round(expected_cost, 10) assert cost != flat_image_cost @@ -171,6 +176,62 @@ def test_gemini_image_edit_cost_uses_output_token_details(): assert cost != all_output_as_image_cost +def test_gemini_image_generation_cost_uses_output_token_details(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + model = "gemini/gemini-3-pro-image-preview" + model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") + + input_text_tokens = 20 + output_text_tokens = 213 + output_image_tokens = 1120 + output_tokens = output_text_tokens + output_image_tokens + image_response = ImageResponse( + data=[ImageObject(b64_json="img1")], + usage=ImageUsage( + input_tokens=input_text_tokens, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=input_text_tokens, + image_tokens=0, + ), + output_tokens=output_tokens, + total_tokens=input_text_tokens + output_tokens, + prompt_tokens=input_text_tokens, + completion_tokens=output_tokens, + prompt_tokens_details={ + "text_tokens": input_text_tokens, + "image_tokens": 0, + }, + completion_tokens_details={ + "text_tokens": output_text_tokens, + "image_tokens": output_image_tokens, + }, + output_tokens_details={ + "text_tokens": output_text_tokens, + "image_tokens": output_image_tokens, + }, + ), + ) + + cost = gemini_image_generation_cost_calculator( + model=model, + image_response=image_response, + ) + + expected_cost = ( + input_text_tokens * model_info["input_cost_per_token"] + + output_text_tokens * model_info["output_cost_per_token"] + + output_image_tokens * model_info["output_cost_per_image_token"] + ) + all_output_as_image_cost = ( + input_text_tokens * model_info["input_cost_per_token"] + + (output_text_tokens + output_image_tokens) + * model_info["output_cost_per_image_token"] + ) + assert round(cost, 10) == round(expected_cost, 10) + assert cost != all_output_as_image_cost + + def test_gemini_image_edit_cost_falls_back_to_flat_image_pricing(): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") From e2ff39793eca0cbeb3256eea7cd6a907b5f84f6e Mon Sep 17 00:00:00 2001 From: Dmitriy Alergant Date: Tue, 19 May 2026 15:50:09 -0400 Subject: [PATCH 09/12] Stop gating Imagen imageSize forwarding --- litellm/llms/gemini/common_utils.py | 15 +++++---------- ...test_gemini_image_generation_transformation.py | 6 +++--- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 5ce223ec95df..d4f7b19d49b5 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -93,8 +93,6 @@ (896, 1280): "3:4", } -IMAGEN_SUPPORTED_IMAGE_SIZES = {"1K", "2K"} - def map_openai_size_to_gemini_image_config( size: str, model: str @@ -111,14 +109,16 @@ def map_openai_size_to_gemini_image_config( if is_gemini_image_model(model): if supports_gemini_image_size(model): image_config["imageSize"] = image_size - elif is_imagen_model(model) and image_size in IMAGEN_SUPPORTED_IMAGE_SIZES: + else: image_config["imageSize"] = image_size return image_config def supports_gemini_image_size(model: str) -> bool: - # Gemini 2.5 Flash image supports aspectRatio but rejects imageSize; newer - # Gemini image models are expected to support both fields. + # gemini-2.5-flash is a legacy model with reduced capability, a one-off + # exception. Newer Nano Banana and Imagen models all support imageSize, and + # newer Gemini image models are widely expected to support it too. Adding a + # model-map feature flag is not justified for this narrow case. return "2.5-flash" not in model @@ -127,11 +127,6 @@ def is_gemini_image_model(model: str) -> bool: return "gemini" in base_model -def is_imagen_model(model: str) -> bool: - base_model = model.split("/", 1)[-1] - return "imagen" in base_model - - def map_openai_image_params_to_gemini( params: Dict[str, Any], model: str, diff --git a/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py b/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py index d2429cfc4ac1..4610d1b99bfa 100644 --- a/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py +++ b/tests/test_litellm/llms/gemini/test_gemini_image_generation_transformation.py @@ -104,7 +104,7 @@ def test_imagen_generation_with_provider_prefix_uses_imagen_params_and_response( assert result.data[0].b64_json == "fake-imagen-image" -def test_imagen_generation_omits_unsupported_openai_size_image_size(): +def test_imagen_generation_forwards_mapped_openai_size_image_size(): config = GoogleImageGenConfig() mapped = config.map_openai_params( @@ -115,7 +115,7 @@ def test_imagen_generation_omits_unsupported_openai_size_image_size(): model="gemini/imagen-4.0-generate-001", drop_params=False, ) - assert mapped == {"aspectRatio": "1:1"} + assert mapped == {"aspectRatio": "1:1", "imageSize": "512"} request = config.transform_image_generation_request( model="gemini/imagen-4.0-generate-001", @@ -127,7 +127,7 @@ def test_imagen_generation_omits_unsupported_openai_size_image_size(): assert request == { "instances": [{"prompt": "Generate a simple app icon"}], - "parameters": {"aspectRatio": "1:1"}, + "parameters": {"aspectRatio": "1:1", "imageSize": "512"}, } From 3df7fb9b59f3e3f0d40f7de82eb1dca3998aea87 Mon Sep 17 00:00:00 2001 From: Dmitriy Alergant Date: Tue, 19 May 2026 15:51:57 -0400 Subject: [PATCH 10/12] Document Gemini image size mapping source --- litellm/llms/gemini/common_utils.py | 6 +++-- .../gemini/image_generation/transformation.py | 2 +- .../llms/gemini/image_usage_transformation.py | 25 +++++++++++-------- 3 files changed, 19 insertions(+), 14 deletions(-) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index d4f7b19d49b5..11ed1aa07444 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -2,7 +2,7 @@ import datetime import json import math -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Sequence, Union import httpx @@ -32,6 +32,8 @@ "21:9": 21 / 9, } +# Supported aspect ratio dimensions from Google Gemini image generation docs: +# https://ai.google.dev/gemini-api/docs/image-generation#aspect_ratios_and_image_size GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO: Dict[tuple[int, int], str] = { (512, 512): "1:1", (1024, 1024): "1:1", @@ -130,7 +132,7 @@ def is_gemini_image_model(model: str) -> bool: def map_openai_image_params_to_gemini( params: Dict[str, Any], model: str, - supported_params: List[str], + supported_params: Sequence[str], optional_params: Optional[Dict[str, Any]] = None, parse_image_config_string: bool = False, ) -> Dict[str, Any]: diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 5c8772c57383..dd72cd809d4a 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, List, Optional import httpx diff --git a/litellm/llms/gemini/image_usage_transformation.py b/litellm/llms/gemini/image_usage_transformation.py index d19e7a3f2e86..5a55bdeffb1c 100644 --- a/litellm/llms/gemini/image_usage_transformation.py +++ b/litellm/llms/gemini/image_usage_transformation.py @@ -1,3 +1,5 @@ +from typing import Any + from litellm.types.utils import ImageUsage, ImageUsageInputTokensDetails @@ -57,14 +59,15 @@ def transform_gemini_image_usage(usage_metadata: dict) -> ImageUsage: if output_tokens > known_output_tokens: output_tokens_details.text_tokens += output_tokens - known_output_tokens - return ImageUsage( - input_tokens=usage_metadata.get("promptTokenCount", 0), - input_tokens_details=input_tokens_details, - output_tokens=output_tokens, - total_tokens=usage_metadata.get("totalTokenCount", 0), - prompt_tokens=usage_metadata.get("promptTokenCount", 0), - prompt_tokens_details=input_tokens_details.model_dump(), - completion_tokens=output_tokens, - completion_tokens_details=output_tokens_details.model_dump(), - output_tokens_details=output_tokens_details.model_dump(), - ) + usage_payload: dict[str, Any] = { + "input_tokens": usage_metadata.get("promptTokenCount", 0), + "input_tokens_details": input_tokens_details, + "output_tokens": output_tokens, + "total_tokens": usage_metadata.get("totalTokenCount", 0), + "prompt_tokens": usage_metadata.get("promptTokenCount", 0), + "prompt_tokens_details": input_tokens_details.model_dump(), + "completion_tokens": output_tokens, + "completion_tokens_details": output_tokens_details.model_dump(), + "output_tokens_details": output_tokens_details.model_dump(), + } + return ImageUsage(**usage_payload) From 61a51a1b29b7eaab9ab2f775ad2868a67df34652 Mon Sep 17 00:00:00 2001 From: DmitriyAlergant Date: Tue, 19 May 2026 19:41:34 -0400 Subject: [PATCH 11/12] chore: retrigger lint --- litellm/llms/gemini/common_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 11ed1aa07444..acf3b05658cc 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -14,7 +14,6 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import TokenCountResponse - GEMINI_IMAGE_ASPECT_RATIOS: Dict[str, float] = { "1:1": 1 / 1, "1:4": 1 / 4, From 8c4a445ca7f632fcd1343ceee2cddbbaafe898d0 Mon Sep 17 00:00:00 2001 From: DmitriyAlergant Date: Tue, 2 Jun 2026 18:53:36 -0400 Subject: [PATCH 12/12] Clarify Gemini candidate count precedence --- litellm/llms/gemini/common_utils.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index acf3b05658cc..e14a17b785a1 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -193,11 +193,13 @@ def get_gemini_image_generation_config( if image_config: generation_config["imageConfig"] = image_config - candidate_count = ( - optional_params.get("candidateCount") - or optional_params.get("candidate_count") - or optional_params.get("sampleCount") - or optional_params.get("n") + candidate_count = next( + ( + optional_params[key] + for key in ("candidateCount", "candidate_count", "sampleCount", "n") + if optional_params.get(key) is not None + ), + None, ) if candidate_count is not None: generation_config["candidateCount"] = candidate_count