gemini image models fixes in /images/generations and /images/edits: usage logging, spend calculation, passing size params

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(

litellm/llms/gemini/common_utils.py

@@ -1,6 +1,8 @@
 import base64
 import datetime
-from typing import Any, Dict, List, Optional, Union
+import json
+import math
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 import httpx
 
@@ -12,6 +14,242 @@
 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,
+    "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,
+}
+
+# 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",
+    (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, 336): "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)
+    }
+    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
+    else:
+        image_config["imageSize"] = image_size
+    return image_config
+
+
+def supports_gemini_image_size(model: str) -> bool:
+    # 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
+
+
+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,
+    supported_params: Sequence[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 is_gemini_image_model(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 = 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
+
+    return generation_config
+
+
+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

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