tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-03 06:23:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add all gemini image models support in image generation (#16526)

Browse Source
This commit is contained in:
Sameer Kankute
2025-11-12 18:26:44 -08:00
committed by GitHub
parent 898f15c33c
commit 394da34a0b
3 changed files with 26 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/completion/image_generation_chat.md
+2 -2
View File
@@ -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
litellm/llms/gemini/image_generation/transformation.py
+11 -6
View File
@@ -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", {})
tests/llm_translation/test_gemini.py
+13 -7
View File
@@ -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)