From 394da34a0beb8be0ae9930dd9caedc963affeff7 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 13 Nov 2025 07:56:44 +0530 Subject: [PATCH] Add all gemini image models support in image generation (#16526) --- .../docs/completion/image_generation_chat.md | 4 ++-- .../gemini/image_generation/transformation.py | 17 ++++++++++------ tests/llm_translation/test_gemini.py | 20 ++++++++++++------- 3 files changed, 26 insertions(+), 15 deletions(-) diff --git a/docs/my-website/docs/completion/image_generation_chat.md b/docs/my-website/docs/completion/image_generation_chat.md index 98b718ef4c..5538b7f8ff 100644 --- a/docs/my-website/docs/completion/image_generation_chat.md +++ b/docs/my-website/docs/completion/image_generation_chat.md @@ -224,8 +224,8 @@ asyncio.run(generate_image()) | Provider | Model | |----------|--------| -| Google AI Studio | `gemini/gemini-2.5-flash-image-preview` | -| Vertex AI | `vertex_ai/gemini-2.5-flash-image-preview` | +| Google AI Studio | `gemini/gemini-2.0-flash-preview-image-generation`, `gemini/gemini-2.5-flash-image-preview` | +| Vertex AI | `vertex_ai/gemini-2.0-flash-preview-image-generation`, `vertex_ai/gemini-2.5-flash-image-preview` | ## Spec diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index f136bd0a40..d47759d0e8 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -21,6 +21,11 @@ else: LiteLLMLoggingObj = Any +FLASH_IMAGE_PREVIEW_MODEL_IDENTIFIERS = ( + "2.0-flash-preview-image", + "2.0-flash-preview-image-generation", + "2.5-flash-image-preview", +) class GoogleImageGenConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta" @@ -97,8 +102,8 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): complete_url = complete_url.rstrip("/") - # Gemini 2.5 Flash Image Preview uses generateContent endpoint - if "2.5-flash-image-preview" in model: + # Gemini Flash Image Preview models use generateContent endpoint + if any(identifier in model for identifier in FLASH_IMAGE_PREVIEW_MODEL_IDENTIFIERS): complete_url = f"{complete_url}/models/{model}:generateContent" else: # All other Imagen models use predict endpoint @@ -152,8 +157,8 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): } } """ - # For Gemini 2.5 Flash Image Preview, use standard Gemini format - if "2.5-flash-image-preview" in model: + # For Gemini Flash Image Preview models, use standard Gemini format + if any(identifier in model for identifier in FLASH_IMAGE_PREVIEW_MODEL_IDENTIFIERS): request_body: dict = { "contents": [ { @@ -212,8 +217,8 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): model_response.data = [] # 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 + if any(identifier in model for identifier in FLASH_IMAGE_PREVIEW_MODEL_IDENTIFIERS): + # Gemini Flash Image Preview models return in candidates format candidates = response_data.get("candidates", []) for candidate in candidates: content = candidate.get("content", {}) diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index f4100d9e8a..b7628a4903 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -290,10 +290,17 @@ def test_gemini_image_generation(): ) -def test_gemini_2_5_flash_image_preview(): +@pytest.mark.parametrize( + "model_name", + [ + "gemini/gemini-2.5-flash-image-preview", + "gemini/gemini-2.0-flash-preview-image-generation", + ], +) +def test_gemini_flash_image_preview_models(model_name: str): """ - 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 + Validate Gemini Flash image preview models route through image_generation() + and invoke the generateContent endpoint returning inline image data. """ from unittest.mock import patch, MagicMock from litellm.types.utils import ImageResponse, ImageObject @@ -321,7 +328,7 @@ def test_gemini_2_5_flash_image_preview(): # Test that the function works without throwing the original 400 error response = litellm.image_generation( - model="gemini/gemini-2.5-flash-image-preview", + model=model_name, prompt="Generate a simple test image", api_key="test_api_key", ) @@ -339,9 +346,9 @@ def test_gemini_2_5_flash_image_preview(): 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) + # Verify it uses generateContent endpoint for Gemini Flash image preview models (not predict) assert ":generateContent" in called_url - assert "gemini-2.5-flash-image-preview" in called_url + assert model_name.split("/", 1)[1] in called_url # Verify request format is Gemini format (not Imagen) request_data = call_args.kwargs.get("json", {}) @@ -356,7 +363,6 @@ def test_gemini_2_5_flash_image_preview(): "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)