fix(cost): use token usage for gemini/vertex image generation when available

- compute image_generation cost from usage token metadata for vertex/gemini\n- map ImageUsage to Usage and reuse generic_cost_per_token\n- fallback to output_cost_per_image when usage metadata missing\n- add tests for token-based path and fallback path

litellm/llms/gemini/image_generation/cost_calculator.py

@@ -2,10 +2,71 @@
 Google AI Image Generation Cost Calculator
 """
 
-from typing import Any
+from typing import Any, Optional
 
 import litellm
-from litellm.types.utils import ImageResponse
+from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
+from litellm.types.utils import (
+    CompletionTokensDetailsWrapper,
+    ImageResponse,
+    PromptTokensDetailsWrapper,
+    Usage,
+)
+
+
+def _calculate_token_based_cost(model: str, image_response: ImageResponse) -> Optional[float]:
+    """
+    Calculate token-based image generation cost when usage metadata is available.
+
+    Falls back to None when usage metadata is missing/incomplete.
+    """
+    usage = image_response.usage
+    if usage is None:
+        return None
+
+    prompt_tokens = usage.input_tokens
+    completion_tokens = usage.output_tokens
+    total_tokens = usage.total_tokens
+
+    if (
+        prompt_tokens is None
+        or completion_tokens is None
+        or total_tokens is None
+    ):
+        return None
+    # ImageResponse may carry a default zeroed usage object even when provider
+    # usage metadata is absent. Treat this as missing usage and fall back.
+    if prompt_tokens == 0 and completion_tokens == 0 and total_tokens == 0:
+        return None
+
+    input_tokens_details = getattr(usage, "input_tokens_details", None)
+    prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
+    if input_tokens_details is not None:
+        prompt_tokens_details = PromptTokensDetailsWrapper(
+            text_tokens=getattr(input_tokens_details, "text_tokens", None),
+            image_tokens=getattr(input_tokens_details, "image_tokens", None),
+            cached_tokens=0,
+        )
+
+    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,
+        ),
+    )
+
+    prompt_cost, completion_cost = generic_cost_per_token(
+        model=model,
+        usage=normalized_usage,
+        custom_llm_provider="gemini",
+    )
+    return prompt_cost + completion_cost
 
 
 def cost_calculator(
@@ -20,6 +81,13 @@ def cost_calculator(
         custom_llm_provider="gemini",
     )
 
+    if isinstance(image_response, ImageResponse):
+        token_based_cost = _calculate_token_based_cost(
+            model=model, image_response=image_response
+        )
+        if token_based_cost is not None:
+            return token_based_cost
+
     output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0
     num_images: int = 0
     if isinstance(image_response, ImageResponse):

litellm/llms/vertex_ai/image_generation/cost_calculator.py

@@ -2,8 +2,71 @@
 Vertex AI Image Generation Cost Calculator
 """
 
+from typing import Optional
+
 import litellm
-from litellm.types.utils import ImageResponse
+from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
+from litellm.types.utils import (
+    CompletionTokensDetailsWrapper,
+    ImageResponse,
+    PromptTokensDetailsWrapper,
+    Usage,
+)
+
+
+def _calculate_token_based_cost(model: str, image_response: ImageResponse) -> Optional[float]:
+    """
+    Calculate token-based image generation cost when usage metadata is available.
+
+    Falls back to None when usage metadata is missing/incomplete.
+    """
+    usage = image_response.usage
+    if usage is None:
+        return None
+
+    prompt_tokens = usage.input_tokens
+    completion_tokens = usage.output_tokens
+    total_tokens = usage.total_tokens
+
+    if (
+        prompt_tokens is None
+        or completion_tokens is None
+        or total_tokens is None
+    ):
+        return None
+    # ImageResponse may carry a default zeroed usage object even when provider
+    # usage metadata is absent. Treat this as missing usage and fall back.
+    if prompt_tokens == 0 and completion_tokens == 0 and total_tokens == 0:
+        return None
+
+    input_tokens_details = getattr(usage, "input_tokens_details", None)
+    prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
+    if input_tokens_details is not None:
+        prompt_tokens_details = PromptTokensDetailsWrapper(
+            text_tokens=getattr(input_tokens_details, "text_tokens", None),
+            image_tokens=getattr(input_tokens_details, "image_tokens", None),
+            cached_tokens=0,
+        )
+
+    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,
+        ),
+    )
+
+    prompt_cost, completion_cost = generic_cost_per_token(
+        model=model,
+        usage=normalized_usage,
+        custom_llm_provider="vertex_ai",
+    )
+    return prompt_cost + completion_cost
 
 
 def cost_calculator(
@@ -18,6 +81,12 @@ def cost_calculator(
         custom_llm_provider="vertex_ai",
     )
 
+    token_based_cost = _calculate_token_based_cost(
+        model=model, image_response=image_response
+    )
+    if token_based_cost is not None:
+        return token_based_cost
+
     output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0
     num_images: int = 0
     if image_response.data:

tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py

@@ -9,9 +9,19 @@ import litellm
 from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
     StandardBuiltInToolCostTracking,
 )
+from litellm.llms.gemini.image_generation.cost_calculator import (
+    cost_calculator as gemini_image_generation_cost_calculator,
+)
+from litellm.llms.vertex_ai.image_generation.cost_calculator import (
+    cost_calculator as vertex_image_generation_cost_calculator,
+)
 from litellm.types.llms.openai import FileSearchTool, WebSearchOptions
 from litellm.types.utils import (
     CompletionTokensDetailsWrapper,
+    ImageObject,
+    ImageResponse,
+    ImageUsage,
+    ImageUsageInputTokensDetails,
     ModelInfo,
     ModelResponse,
     PromptTokensDetailsWrapper,
@@ -837,6 +847,140 @@ def test_gemini_image_generation_cost_with_zero_text_tokens(model: str):
     )
 
 
+def test_vertex_image_generation_cost_prefers_token_usage_metadata():
+    """
+    When usage metadata exists on image responses, Vertex image generation cost
+    should be calculated from token pricing, not flat output_cost_per_image.
+    """
+    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
+    litellm.model_cost = litellm.get_model_cost_map(url="")
+
+    model = "gemini-3.1-flash-image-preview"
+    model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai")
+
+    input_text_tokens = 50
+    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 = vertex_image_generation_cost_calculator(
+        model=model,
+        image_response=image_response,
+    )
+
+    expected_prompt_cost = prompt_tokens * model_info["input_cost_per_token"]
+    expected_completion_cost = output_image_tokens * model_info["output_cost_per_image_token"]
+    expected_total_cost = expected_prompt_cost + expected_completion_cost
+
+    assert round(cost, 10) == round(expected_total_cost, 10)
+    # Ensure this is not falling back to flat per-image pricing.
+    assert cost != len(image_response.data) * model_info["output_cost_per_image"]
+
+
+def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing():
+    """
+    Without usage metadata, Vertex image generation cost should fall back to
+    output_cost_per_image * number_of_images.
+    """
+    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
+    litellm.model_cost = litellm.get_model_cost_map(url="")
+
+    model = "gemini-3.1-flash-image-preview"
+    model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai")
+
+    image_response = ImageResponse(
+        data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]
+    )
+
+    cost = vertex_image_generation_cost_calculator(
+        model=model,
+        image_response=image_response,
+    )
+
+    expected_cost = len(image_response.data) * model_info["output_cost_per_image"]
+    assert round(cost, 10) == round(expected_cost, 10)
+
+
+def test_gemini_image_generation_cost_prefers_token_usage_metadata():
+    """
+    When usage metadata exists on image responses, Gemini image generation cost
+    should be calculated from token pricing, not flat output_cost_per_image.
+    """
+    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_generation_cost_calculator(
+        model=model,
+        image_response=image_response,
+    )
+
+    expected_prompt_cost = prompt_tokens * model_info["input_cost_per_token"]
+    expected_completion_cost = output_image_tokens * model_info["output_cost_per_image_token"]
+    expected_total_cost = expected_prompt_cost + expected_completion_cost
+
+    assert round(cost, 10) == round(expected_total_cost, 10)
+    # Ensure this is not falling back to flat per-image pricing.
+    assert cost != len(image_response.data) * model_info["output_cost_per_image"]
+
+
+def test_gemini_image_generation_cost_falls_back_to_flat_image_pricing():
+    """
+    Without usage metadata, Gemini image generation cost should fall back to
+    output_cost_per_image * number_of_images.
+    """
+    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_generation_cost_calculator(
+        model=model,
+        image_response=image_response,
+    )
+
+    expected_cost = len(image_response.data) * model_info["output_cost_per_image"]
+    assert round(cost, 10) == round(expected_cost, 10)
+
+
 def test_bedrock_anthropic_prompt_caching():
     """Test Bedrock Anthropic models with prompt caching return correct costs."""
     model = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"