[Feat] Add Google AI Studio Imagen4 model family (#13065)

* add gemini

* add init files

* add get_gemini_image_generation_config

* refactor transform

* TestGoogleImageGen

* fix transform

* fix transform

* add gemini_image_cost_calculator

* add cost tracking for gemini/imagen models

* docs image gen

* docs image gen

* test_get_model_info_gemini

docs/my-website/docs/image_generation.md

@@ -278,3 +278,16 @@ response = litellm.image_generation(
 )
 print(f"response: {response}")
 ```
+
+## Supported Providers
+
+| Provider | Documentation Link |
+|----------|-------------------|
+| OpenAI | [OpenAI Image Generation →](./providers/openai) |
+| Azure OpenAI | [Azure OpenAI Image Generation →](./providers/azure/azure) |
+| Google AI Studio | [Google AI Studio Image Generation →](./providers/google_ai_studio/image_gen) |
+| Vertex AI | [Vertex AI Image Generation →](./providers/vertex_image) |
+| AWS Bedrock | [Bedrock Image Generation →](./providers/bedrock) |
+| Recraft | [Recraft Image Generation →](./providers/recraft#image-generation) |
+| Xinference | [Xinference Image Generation →](./providers/xinference#image-generation) |
+| Nscale | [Nscale Image Generation →](./providers/nscale#image-generation) | 

docs/my-website/docs/providers/google_ai_studio/image_gen.md

@@ -0,0 +1,214 @@
+import Tabs from '@theme/Tabs';
+import TabItem from '@theme/TabItem';
+
+# Google AI Studio Image Generation
+
+Google AI Studio provides powerful image generation capabilities using Google's Imagen models to create high-quality images from text descriptions.
+
+## Overview
+
+| Property | Details |
+|----------|---------|
+| Description | Google AI Studio Image Generation uses Google's Imagen models to generate high-quality images from text descriptions. |
+| Provider Route on LiteLLM | `gemini/` |
+| Provider Doc | [Google AI Studio Image Generation ↗](https://ai.google.dev/gemini-api/docs/imagen) |
+| Supported Operations | [`/images/generations`](#image-generation) |
+
+## Setup
+
+### API Key
+
+```python showLineNumbers
+# Set your Google AI Studio API key
+import os
+os.environ["GEMINI_API_KEY"] = "your-api-key-here"
+```
+
+Get your API key from [Google AI Studio](https://aistudio.google.com/app/apikey).
+
+## Image Generation
+
+### Usage - LiteLLM Python SDK
+
+<Tabs>
+<TabItem value="basic" label="Basic Usage">
+
+```python showLineNumbers title="Basic Image Generation"
+import litellm
+import os
+
+# Set your API key
+os.environ["GEMINI_API_KEY"] = "your-api-key-here"
+
+# Generate a single image
+response = litellm.image_generation(
+    model="gemini/imagen-4.0-generate-preview-06-06",
+    prompt="A cute baby sea otter swimming in crystal clear water"
+)
+
+print(response.data[0].url)
+```
+
+</TabItem>
+
+<TabItem value="async" label="Async Usage">
+
+```python showLineNumbers title="Async Image Generation"
+import litellm
+import asyncio
+import os
+
+async def generate_image():
+    # Set your API key
+    os.environ["GEMINI_API_KEY"] = "your-api-key-here"
+    
+    # Generate image asynchronously
+    response = await litellm.aimage_generation(
+        model="gemini/imagen-4.0-generate-preview-06-06",
+        prompt="A beautiful sunset over mountains with vibrant colors",
+        n=1,
+    )
+    
+    print(response.data[0].url)
+    return response
+
+# Run the async function
+asyncio.run(generate_image())
+```
+
+</TabItem>
+
+<TabItem value="advanced" label="Advanced Parameters">
+
+```python showLineNumbers title="Advanced Image Generation with Parameters"
+import litellm
+import os
+
+# Set your API key
+os.environ["GEMINI_API_KEY"] = "your-api-key-here"
+
+# Generate image with additional parameters
+response = litellm.image_generation(
+    model="gemini/imagen-4.0-generate-preview-06-06",
+    prompt="A futuristic cityscape at night with neon lights",
+    n=1,
+    size="1024x1024",
+    quality="standard",
+    response_format="url"
+)
+
+for image in response.data:
+    print(f"Generated image URL: {image.url}")
+```
+
+</TabItem>
+</Tabs>
+
+### Usage - LiteLLM Proxy Server
+
+#### 1. Configure your config.yaml
+
+```yaml showLineNumbers title="Google AI Studio Image Generation Configuration"
+model_list:
+  - model_name: google-imagen
+    litellm_params:
+      model: gemini/imagen-4.0-generate-preview-06-06
+      api_key: os.environ/GEMINI_API_KEY
+  model_info:
+    mode: image_generation
+
+general_settings:
+  master_key: sk-1234
+```
+
+#### 2. Start LiteLLM Proxy Server
+
+```bash showLineNumbers title="Start LiteLLM Proxy Server"
+litellm --config /path/to/config.yaml
+
+# RUNNING on http://0.0.0.0:4000
+```
+
+#### 3. Make requests with OpenAI Python SDK
+
+<Tabs>
+<TabItem value="openai-sdk" label="OpenAI SDK">
+
+```python showLineNumbers title="Google AI Studio Image Generation via Proxy - OpenAI SDK"
+from openai import OpenAI
+
+# Initialize client with your proxy URL
+client = OpenAI(
+    base_url="http://localhost:4000",  # Your proxy URL
+    api_key="sk-1234"                  # Your proxy API key
+)
+
+# Generate image
+response = client.images.generate(
+    model="google-imagen",
+    prompt="A majestic eagle soaring over snow-capped mountains",
+    n=1,
+    size="1024x1024"
+)
+
+print(response.data[0].url)
+```
+
+</TabItem>
+
+<TabItem value="litellm-sdk" label="LiteLLM SDK">
+
+```python showLineNumbers title="Google AI Studio Image Generation via Proxy - LiteLLM SDK"
+import litellm
+
+# Configure LiteLLM to use your proxy
+response = litellm.image_generation(
+    model="litellm_proxy/google-imagen",
+    prompt="A serene Japanese garden with cherry blossoms",
+    api_base="http://localhost:4000",
+    api_key="sk-1234"
+)
+
+print(response.data[0].url)
+```
+
+</TabItem>
+
+<TabItem value="curl" label="cURL">
+
+```bash showLineNumbers title="Google AI Studio Image Generation via Proxy - cURL"
+curl --location 'http://localhost:4000/v1/images/generations' \
+--header 'Content-Type: application/json' \
+--header 'Authorization: Bearer sk-1234' \
+--data '{
+    "model": "google-imagen",
+    "prompt": "A cozy coffee shop interior with warm lighting",
+    "n": 1,
+    "size": "1024x1024"
+}'
+```
+
+</TabItem>
+</Tabs>
+
+## Supported Parameters
+
+Google AI Studio Image Generation supports the following OpenAI-compatible parameters:
+
+| Parameter | Type | Description | Default | Example |
+|-----------|------|-------------|---------|---------|
+| `prompt` | string | Text description of the image to generate | Required | `"A sunset over the ocean"` |
+| `model` | string | The model to use for generation | Required | `"gemini/imagen-4.0-generate-preview-06-06"` |
+| `n` | integer | Number of images to generate (1-4) | `1` | `2` |
+| `size` | string | Image dimensions | `"1024x1024"` | `"512x512"`, `"1024x1024"` |
+
+1. Create an account at [Google AI Studio](https://aistudio.google.com/)
+2. Generate an API key from [API Keys section](https://aistudio.google.com/app/apikey)
+3. Set your `GEMINI_API_KEY` environment variable
+4. Start generating images using LiteLLM
+
+## Additional Resources
+
+- [Google AI Studio Documentation](https://ai.google.dev/gemini-api/docs)
+- [Imagen Model Overview](https://ai.google.dev/gemini-api/docs/imagen)
+- [LiteLLM Image Generation Guide](../../completion/image_generation)

litellm/cost_calculator.py

@@ -793,6 +793,14 @@ def completion_cost(  # noqa: PLR0915
                             model=model,
                             image_response=completion_response,
                         )
+                    elif custom_llm_provider == litellm.LlmProviders.GEMINI.value:
+                        from litellm.llms.gemini.image_generation.cost_calculator import (
+                            cost_calculator as gemini_image_cost_calculator,
+                        )
+                        return gemini_image_cost_calculator(
+                            model=model,
+                            image_response=completion_response,
+                        )
                     else:
                         return default_image_cost_calculator(
                             model=model,

litellm/images/main.py

@@ -411,6 +411,8 @@ def image_generation(  # noqa: PLR0915
         #########################################################
         elif custom_llm_provider in (
             litellm.LlmProviders.RECRAFT,
+            litellm.LlmProviders.GEMINI,
+            
         ):
             if image_generation_config is None:
                 raise ValueError(f"image generation config is not supported for {custom_llm_provider}")

litellm/llms/gemini/image_generation/__init__.py

@@ -0,0 +1,13 @@
+from litellm.llms.base_llm.image_generation.transformation import (
+    BaseImageGenerationConfig,
+)
+
+from .transformation import GoogleImageGenConfig
+
+__all__ = [
+    "GoogleImageGenConfig",
+]
+
+
+def get_gemini_image_generation_config(model: str) -> BaseImageGenerationConfig:
+    return GoogleImageGenConfig()

litellm/llms/gemini/image_generation/cost_calculator.py

@@ -0,0 +1,30 @@
+"""
+Google AI Image Generation Cost Calculator
+"""
+
+from typing import Any
+
+import litellm
+from litellm.types.utils import ImageResponse
+
+
+def cost_calculator(
+    model: str,
+    image_response: Any,
+) -> float:
+    """
+    Vertex AI Image Generation Cost Calculator
+    """
+    _model_info = litellm.get_model_info(
+        model=model,
+        custom_llm_provider="gemini",
+    )
+
+    output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0
+    num_images: int = 0
+    if isinstance(image_response, ImageResponse):
+        if image_response.data:
+            num_images = len(image_response.data)
+        return output_cost_per_image * num_images
+    else:
+        raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}")

litellm/llms/gemini/image_generation/transformation.py

@@ -0,0 +1,200 @@
+from typing import TYPE_CHECKING, Any, List, Optional
+
+import httpx
+
+from litellm.llms.base_llm.image_generation.transformation import (
+    BaseImageGenerationConfig,
+)
+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
+
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
+
+    LiteLLMLoggingObj = _LiteLLMLoggingObj
+else:
+    LiteLLMLoggingObj = Any
+
+
+class GoogleImageGenConfig(BaseImageGenerationConfig):
+    DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta"
+    
+    def get_supported_openai_params(
+        self, model: str
+    ) -> List[OpenAIImageGenerationOptionalParams]:
+        """
+        Google AI Imagen API supported parameters
+        https://ai.google.dev/gemini-api/docs/imagen
+        """
+        return [
+            "n",
+            "size"
+        ]
+    
+    def map_openai_params(
+        self,
+        non_default_params: dict,
+        optional_params: dict,
+        model: str,
+        drop_params: bool,
+    ) -> dict:
+        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
+    
+
+    def _map_size_to_aspect_ratio(self, size: str) -> str:
+        """
+        https://ai.google.dev/gemini-api/docs/image-generation
+
+        """
+        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 get_complete_url(
+        self,
+        api_base: Optional[str],
+        api_key: Optional[str],
+        model: str,
+        optional_params: dict,
+        litellm_params: dict,
+        stream: Optional[bool] = None,
+    ) -> str:
+        """
+        Get the complete url for the request
+        
+        Google AI API format: https://generativelanguage.googleapis.com/v1beta/models/{model}:predict
+        """
+        complete_url: str = (
+            api_base 
+            or get_secret_str("GEMINI_API_BASE") 
+            or self.DEFAULT_BASE_URL
+        )
+
+        complete_url = complete_url.rstrip("/")
+        complete_url = f"{complete_url}/models/{model}:predict"
+        return complete_url
+
+    def validate_environment(
+        self,
+        headers: dict,
+        model: str,
+        messages: List[AllMessageValues],
+        optional_params: dict,
+        litellm_params: dict,
+        api_key: Optional[str] = None,
+        api_base: Optional[str] = None,
+    ) -> dict:
+        final_api_key: Optional[str] = (
+            api_key or 
+            get_secret_str("GEMINI_API_KEY")
+        )
+        if not final_api_key:
+            raise ValueError("GEMINI_API_KEY is not set")
+        
+        headers["x-goog-api-key"] = final_api_key
+        headers["Content-Type"] = "application/json"
+        return headers
+
+    def transform_image_generation_request(
+        self,
+        model: str,
+        prompt: str,
+        optional_params: dict,
+        litellm_params: dict,
+        headers: dict,
+    ) -> dict:
+        """
+        Transform the image generation request to Google AI Imagen format
+        
+        Google AI API format:
+        {
+          "instances": [
+            {
+              "prompt": "Robot holding a red skateboard"
+            }
+          ],
+          "parameters": {
+            "sampleCount": 4,
+            "aspectRatio": "1:1",
+            "personGeneration": "allow_adult"
+          }
+        }
+        """
+        from litellm.types.llms.gemini import (
+            GeminiImageGenerationInstance,
+            GeminiImageGenerationParameters,
+        )
+        request_body: GeminiImageGenerationRequest = GeminiImageGenerationRequest(
+            instances=[
+                GeminiImageGenerationInstance(
+                    prompt=prompt
+                )
+            ],
+            parameters=GeminiImageGenerationParameters(**optional_params)
+        )
+        return request_body.model_dump(exclude_none=True)
+
+    def transform_image_generation_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        model_response: ImageResponse,
+        logging_obj: LiteLLMLoggingObj,
+        request_data: dict,
+        optional_params: dict,
+        litellm_params: dict,
+        encoding: Any,
+        api_key: Optional[str] = None,
+        json_mode: Optional[bool] = None,
+    ) -> ImageResponse:
+        """
+        Transform Google AI Imagen response to litellm ImageResponse format
+        """
+        try:
+            response_data = raw_response.json()
+        except Exception as e:
+            raise self.get_error_class(
+                error_message=f"Error transforming image generation response: {e}",
+                status_code=raw_response.status_code,
+                headers=raw_response.headers,
+            )
+        
+        if not model_response.data:
+            model_response.data = []
+        
+        # Google AI returns predictions with generated images
+        predictions = response_data.get("predictions", [])
+        for prediction in predictions:
+            # Google AI returns base64 encoded images in the prediction
+            generated_images = prediction.get("generatedImages", [])
+            for image_data in generated_images:
+                model_response.data.append(ImageObject(
+                    b64_json=image_data.get("bytesBase64Encoded", None),
+                    url=None,  # Google AI returns base64, not URLs
+                ))
+        
+        return model_response

litellm/llms/recraft/cost_calculator.py

@@ -1,10 +1,12 @@
+from typing import Any
+
 import litellm
 from litellm.types.utils import ImageResponse
 
 
 def cost_calculator(
     model: str,
-    image_response: ImageResponse,
+    image_response: Any,
 ) -> float:
     """
     Recraft image generation cost calculator
@@ -15,6 +17,9 @@ def cost_calculator(
     )
     output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0
     num_images: int = 0
-    if image_response.data:
-        num_images = len(image_response.data)
-    return output_cost_per_image * num_images
+    if isinstance(image_response, ImageResponse):
+        if image_response.data:
+            num_images = len(image_response.data)
+        return output_cost_per_image * num_images
+    else:
+        raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}")

litellm/model_prices_and_context_window_backup.json

@@ -10094,6 +10094,42 @@
         "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models",
         "supports_tool_choice": true
     },
+    "gemini/imagen-4.0-generate-preview-06-06": {
+        "output_cost_per_image": 0.04,
+        "litellm_provider": "gemini",
+        "mode": "image_generation",
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
+    },
+    "gemini/imagen-4.0-ultra-generate-preview-06-06": {
+        "output_cost_per_image": 0.06,
+        "litellm_provider": "gemini",
+        "mode": "image_generation",
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
+    },
+    "gemini/imagen-4.0-fast-generate-preview-06-06": {
+        "output_cost_per_image": 0.02,
+        "litellm_provider": "gemini",
+        "mode": "image_generation",
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
+    },
+    "gemini/imagen-3.0-generate-002": {
+        "output_cost_per_image": 0.04,
+        "litellm_provider": "gemini",
+        "mode": "image_generation",
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
+    },
+    "gemini/imagen-3.0-generate-001": {
+        "output_cost_per_image": 0.04,
+        "litellm_provider": "gemini",
+        "mode": "image_generation",
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
+    },
+    "gemini/imagen-3.0-fast-generate-001": {
+        "output_cost_per_image": 0.02,
+        "litellm_provider": "gemini",
+        "mode": "image_generation",
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
+    },
     "command-a-03-2025": {
         "max_tokens": 8000,
         "max_input_tokens": 256000,

litellm/types/llms/gemini.py

@@ -150,3 +150,74 @@ class BidiGenerateContentSetup(TypedDict, total=False):
 
     outputAudioTranscription: dict
     """The output audio transcription to be used for the realtime session."""
+
+
+# Image Generation Types
+from pydantic import BaseModel
+
+
+class GeminiImageGenerationInstance(TypedDict):
+    """Instance data for Gemini image generation request"""
+    prompt: str
+
+
+class GeminiImageGenerationParameters(BaseModel):
+    """Parameters for Gemini image generation request"""
+    sampleCount: Optional[int] = None
+    """Number of images to generate (maps to OpenAI 'n' parameter)"""
+    
+    aspectRatio: Optional[str] = None
+    """Aspect ratio for generated images (e.g., '1:1', '16:9', '9:16', '4:3', '3:4')"""
+    
+    personGeneration: Optional[str] = None
+    """Controls person generation in images"""
+    
+    # Additional parameters that might be passed through
+    background: Optional[str] = None
+    """Background specification"""
+    
+    input_fidelity: Optional[str] = None
+    """Input fidelity specification"""
+    
+    moderation: Optional[str] = None
+    """Moderation settings"""
+    
+    output_compression: Optional[str] = None
+    """Output compression settings"""
+    
+    output_format: Optional[str] = None
+    """Output format specification"""
+    
+    quality: Optional[str] = None
+    """Quality settings"""
+    
+    response_format: Optional[str] = None
+    """Response format specification"""
+    
+    style: Optional[str] = None
+    """Style specification"""
+    
+    user: Optional[str] = None
+    """User specification"""
+
+
+class GeminiImageGenerationRequest(BaseModel):
+    """Complete request body for Gemini image generation"""
+    instances: List[GeminiImageGenerationInstance]
+    parameters: GeminiImageGenerationParameters
+
+
+class GeminiGeneratedImage(TypedDict):
+    """Individual generated image data from Gemini response"""
+    bytesBase64Encoded: str
+    """Base64 encoded image data"""
+
+
+class GeminiImageGenerationPrediction(TypedDict):
+    """Prediction object containing generated images"""
+    generatedImages: List[GeminiGeneratedImage]
+
+
+class GeminiImageGenerationResponse(TypedDict):
+    """Complete response body from Gemini image generation API"""
+    predictions: List[GeminiImageGenerationPrediction]

litellm/utils.py

@@ -7187,6 +7187,12 @@ class ProviderConfigManager:
             )
 
             return get_recraft_image_generation_config(model)
+        elif LlmProviders.GEMINI == provider:
+            from litellm.llms.gemini.image_generation import (
+                get_gemini_image_generation_config,
+            )
+
+            return get_gemini_image_generation_config(model)
         return None
 
     @staticmethod

model_prices_and_context_window.json

@@ -10094,6 +10094,42 @@
         "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models",
         "supports_tool_choice": true
     },
+    "gemini/imagen-4.0-generate-preview-06-06": {
+        "output_cost_per_image": 0.04,
+        "litellm_provider": "gemini",
+        "mode": "image_generation",
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
+    },
+    "gemini/imagen-4.0-ultra-generate-preview-06-06": {
+        "output_cost_per_image": 0.06,
+        "litellm_provider": "gemini",
+        "mode": "image_generation",
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
+    },
+    "gemini/imagen-4.0-fast-generate-preview-06-06": {
+        "output_cost_per_image": 0.02,
+        "litellm_provider": "gemini",
+        "mode": "image_generation",
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
+    },
+    "gemini/imagen-3.0-generate-002": {
+        "output_cost_per_image": 0.04,
+        "litellm_provider": "gemini",
+        "mode": "image_generation",
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
+    },
+    "gemini/imagen-3.0-generate-001": {
+        "output_cost_per_image": 0.04,
+        "litellm_provider": "gemini",
+        "mode": "image_generation",
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
+    },
+    "gemini/imagen-3.0-fast-generate-001": {
+        "output_cost_per_image": 0.02,
+        "litellm_provider": "gemini",
+        "mode": "image_generation",
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
+    },
     "command-a-03-2025": {
         "max_tokens": 8000,
         "max_input_tokens": 256000,

tests/image_gen_tests/test_image_generation.py

@@ -169,6 +169,9 @@ class TestRecraftImageGeneration(BaseImageGenTest):
     def get_base_image_generation_call_args(self) -> dict:
         return {"model": "recraft/recraftv3"}
 
+class TestGoogleImageGen(BaseImageGenTest):
+    def get_base_image_generation_call_args(self) -> dict:
+        return {"model": "gemini/imagen-4.0-generate-preview-06-06"}
 
 class TestAzureOpenAIDalle3(BaseImageGenTest):
     def get_base_image_generation_call_args(self) -> dict:

tests/test_litellm/test_utils.py

@@ -11,6 +11,7 @@ sys.path.insert(
 )  # Adds the parent directory to the system path
 
 import litellm
+from litellm.proxy.utils import is_valid_api_key
 from litellm.types.utils import (
     Delta,
     LlmProviders,
@@ -23,7 +24,6 @@ from litellm.utils import (
     get_llm_provider,
     get_optional_params_image_gen,
 )
-from litellm.proxy.utils import is_valid_api_key
 
 # Adds the parent directory to the system path
 
@@ -558,6 +558,7 @@ def test_get_model_info_gemini():
             model.startswith("gemini/")
             and not "gemma" in model
             and not "learnlm" in model
+            and not "imagen" in model
         ):
             assert info.get("tpm") is not None, f"{model} does not have tpm"
             assert info.get("rpm") is not None, f"{model} does not have rpm"
@@ -2158,6 +2159,7 @@ def test_image_response_utils():
 
 def test_is_valid_api_key():
     import hashlib
+
     # Valid sk- keys
     assert is_valid_api_key("sk-abc123")
     assert is_valid_api_key("sk-ABC_123-xyz")