Merge pull request #14715 from timelfrink/fix/issue-14120-gemini-2.5-flash-image-preview

Fix: gemini-2.5-flash-image-preview model routing for image generation

litellm/llms/gemini/image_generation/transformation.py

@@ -85,17 +85,25 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
     ) -> str:
         """
         Get the complete url for the request
-        
-        Google AI API format: https://generativelanguage.googleapis.com/v1beta/models/{model}:predict
+
+        Gemini 2.5 Flash Image Preview: :generateContent
+        Other Imagen models: :predict
         """
         complete_url: str = (
-            api_base 
-            or get_secret_str("GEMINI_API_BASE") 
+            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"
+
+        # Gemini 2.5 Flash Image Preview uses generateContent endpoint
+        if "2.5-flash-image-preview" in model:
+            complete_url = f"{complete_url}/models/{model}:generateContent"
+        else:
+            # All other Imagen models use predict endpoint
+            complete_url = f"{complete_url}/models/{model}:predict"
+
         return complete_url
 
     def validate_environment(
@@ -128,35 +136,52 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
         headers: dict,
     ) -> dict:
         """
-        Transform the image generation request to Google AI Imagen format
-        
-        Google AI API format:
+        Transform the image generation request to Gemini format
+
+        For Gemini 2.5 Flash Image Preview, use the standard Gemini format with response_modalities:
         {
-          "instances": [
+          "contents": [
             {
-              "prompt": "Robot holding a red skateboard"
+              "parts": [
+                {"text": "Generate an image of..."}
+              ]
             }
           ],
-          "parameters": {
-            "sampleCount": 4,
-            "aspectRatio": "1:1",
-            "personGeneration": "allow_adult"
+          "generationConfig": {
+            "response_modalities": ["IMAGE", "TEXT"]
           }
         }
         """
-        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)
+        # For Gemini 2.5 Flash Image Preview, use standard Gemini format
+        if "2.5-flash-image-preview" in model:
+            request_body: dict = {
+                "contents": [
+                    {
+                        "parts": [
+                            {"text": prompt}
+                        ]
+                    }
+                ],
+                "generationConfig": {
+                    "response_modalities": ["IMAGE", "TEXT"]
+                }
+            }
+            return request_body
+        else:
+            # For other Imagen models, use the original Imagen format
+            from litellm.types.llms.gemini import (
+                GeminiImageGenerationInstance,
+                GeminiImageGenerationParameters,
+            )
+            request_body_obj: GeminiImageGenerationRequest = GeminiImageGenerationRequest(
+                instances=[
+                    GeminiImageGenerationInstance(
+                        prompt=prompt
+                    )
+                ],
+                parameters=GeminiImageGenerationParameters(**optional_params)
+            )
+            return request_body_obj.model_dump(exclude_none=True)
 
     def transform_image_generation_response(
         self,
@@ -185,14 +210,30 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
         
         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
-            model_response.data.append(ImageObject(
-                b64_json=prediction.get("bytesBase64Encoded", None),
-                url=None,  # Google AI returns base64, not URLs
-            ))
-        
+
+        # Handle different response formats based on model
+        if "2.5-flash-image-preview" in model:
+            # Gemini 2.5 Flash Image Preview returns in candidates format
+            candidates = response_data.get("candidates", [])
+            for candidate in candidates:
+                content = candidate.get("content", {})
+                parts = content.get("parts", [])
+                for part in parts:
+                    # Look for inlineData with image
+                    if "inlineData" in part:
+                        inline_data = part["inlineData"]
+                        if "data" in inline_data:
+                            model_response.data.append(ImageObject(
+                                b64_json=inline_data["data"],
+                                url=None,
+                            ))
+        else:
+            # Original Imagen format - predictions with generated images
+            predictions = response_data.get("predictions", [])
+            for prediction in predictions:
+                # Google AI returns base64 encoded images in the prediction
+                model_response.data.append(ImageObject(
+                    b64_json=prediction.get("bytesBase64Encoded", None),
+                    url=None,  # Google AI returns base64, not URLs
+                ))
         return model_response

litellm/model_prices_and_context_window_backup.json

@@ -9126,7 +9126,7 @@
         "max_tokens": 65535,
         "max_video_length": 1,
         "max_videos_per_prompt": 10,
-        "mode": "chat",
+        "mode": "image_generation",
         "output_cost_per_image": 0.039,
         "output_cost_per_reasoning_token": 3e-05,
         "output_cost_per_token": 3e-05,
@@ -10489,7 +10489,7 @@
         "max_tokens": 65535,
         "max_video_length": 1,
         "max_videos_per_prompt": 10,
-        "mode": "chat",
+        "mode": "image_generation",
         "output_cost_per_image": 0.039,
         "output_cost_per_reasoning_token": 3e-05,
         "output_cost_per_token": 3e-05,

model_prices_and_context_window.json

@@ -9126,7 +9126,7 @@
         "max_tokens": 65535,
         "max_video_length": 1,
         "max_videos_per_prompt": 10,
-        "mode": "chat",
+        "mode": "image_generation",
         "output_cost_per_image": 0.039,
         "output_cost_per_reasoning_token": 3e-05,
         "output_cost_per_token": 3e-05,
@@ -10489,7 +10489,7 @@
         "max_tokens": 65535,
         "max_video_length": 1,
         "max_videos_per_prompt": 10,
-        "mode": "chat",
+        "mode": "image_generation",
         "output_cost_per_image": 0.039,
         "output_cost_per_reasoning_token": 3e-05,
         "output_cost_per_token": 3e-05,

tests/llm_translation/test_gemini.py

@@ -272,6 +272,119 @@ def test_gemini_image_generation():
     assert response.choices[0].message.images[0]["image_url"]["url"].startswith("data:image/png;base64,")
 
 
+def test_gemini_2_5_flash_image_preview():
+    """
+    Test for GitHub issue #14120 - gemini-2.5-flash-image-preview model routing fix
+    Validates that the model correctly routes to image generation instead of chat completion
+    """
+    from unittest.mock import patch, MagicMock
+    from litellm.types.utils import ImageResponse, ImageObject
+
+    # Mock successful response to avoid API limits
+    mock_response = ImageResponse()
+    mock_response.data = [ImageObject(b64_json="test_base64_data", url=None)]
+
+    with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post:
+        # Mock successful HTTP response
+        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
+
+        # Test that the function works without throwing the original 400 error
+        response = litellm.image_generation(
+            model="gemini/gemini-2.5-flash-image-preview",
+            prompt="Generate a simple test image",
+            api_key="test_api_key"
+        )
+
+        # Validate response structure
+        assert response is not None
+        assert hasattr(response, 'data')
+        assert response.data is not None
+        assert len(response.data) > 0
+
+        # Validate the correct endpoint was called
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        called_url = call_args[0][0] if call_args[0] else call_args.kwargs.get('url', '')
+
+        # Verify it uses generateContent endpoint for gemini-2.5-flash-image-preview (not predict)
+        assert ":generateContent" in called_url
+        assert "gemini-2.5-flash-image-preview" in called_url
+
+        # Verify request format is Gemini format (not Imagen)
+        request_data = call_args.kwargs.get('json', {})
+        assert "contents" in request_data
+        assert "parts" in request_data["contents"][0]
+
+        # Verify response_modalities is set correctly for image generation
+        assert "generationConfig" in request_data
+        assert "response_modalities" in request_data["generationConfig"]
+        assert request_data["generationConfig"]["response_modalities"] == ["IMAGE", "TEXT"]
+
+
+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)
+    """
+    from unittest.mock import patch, MagicMock
+    from litellm.types.utils import ImageResponse, ImageObject
+
+    with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post:
+        # Mock successful HTTP response for Imagen
+        mock_http_response = MagicMock()
+        mock_http_response.json.return_value = {
+            "predictions": [
+                {
+                    "bytesBase64Encoded": "test_base64_image_data"
+                }
+            ]
+        }
+        mock_http_response.status_code = 200
+        mock_post.return_value = mock_http_response
+
+        # Test an Imagen model
+        response = litellm.image_generation(
+            model="gemini/imagen-3.0-generate-001",
+            prompt="Generate a simple test image",
+            api_key="test_api_key"
+        )
+
+        # Validate response structure
+        assert response is not None
+        assert hasattr(response, 'data')
+
+        # Validate the correct endpoint was called for Imagen models
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        called_url = call_args[0][0] if call_args[0] else call_args.kwargs.get('url', '')
+
+        # Verify Imagen models use predict endpoint (not generateContent)
+        assert ":predict" in called_url
+        assert "imagen-3.0-generate-001" in called_url
+        assert ":generateContent" not in called_url
+
+        # Verify request format is Imagen format (not Gemini)
+        request_data = call_args.kwargs.get('json', {})
+        assert "instances" in request_data
+        assert "parameters" in request_data
+
+
 def test_gemini_thinking():
     litellm._turn_on_debug()
     from litellm.types.utils import Message, CallTypes