diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index e57364fd28..f136bd0a40 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/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 \ No newline at end of file diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 29100016bb..43ea3af320 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/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, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 29100016bb..43ea3af320 100644 --- a/model_prices_and_context_window.json +++ b/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, diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index 8fcc85aba8..120ffba960 100644 --- a/tests/llm_translation/test_gemini.py +++ b/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