scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

Add better handling image generation for gemini models

Browse source
This commit is contained in:
Sameer Kankute 2025-12-01 10:42:38 +05:30
parent b949ec90db
commit 02510a908f
1 changed files with 3 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/llms/gemini/image_generation/transformation.py
View file

@ -21,12 +21,6 @@ else:
LiteLLMLoggingObj = Any
FLASH_IMAGE_PREVIEW_MODEL_IDENTIFIERS = (
"2.0-flash-preview-image",
"2.0-flash-preview-image-generation",
"2.5-flash-image-preview",
"3-pro-image-preview",
)
class GoogleImageGenConfig(BaseImageGenerationConfig):
DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta"
@ -104,7 +98,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
complete_url = complete_url.rstrip("/")
# Gemini Flash Image Preview models use generateContent endpoint
if any(identifier in model for identifier in FLASH_IMAGE_PREVIEW_MODEL_IDENTIFIERS):
if "gemini" in model:
complete_url = f"{complete_url}/models/{model}:generateContent"
else:
# All other Imagen models use predict endpoint
@ -159,7 +153,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
}
"""
# For Gemini Flash Image Preview models, use standard Gemini format
if any(identifier in model for identifier in FLASH_IMAGE_PREVIEW_MODEL_IDENTIFIERS):
if "gemini" in model:
request_body: dict = {
"contents": [
{
@ -218,7 +212,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
model_response.data = []
# Handle different response formats based on model
if any(identifier in model for identifier in FLASH_IMAGE_PREVIEW_MODEL_IDENTIFIERS):
if "gemini" in model:
# Gemini Flash Image Preview models return in candidates format
candidates = response_data.get("candidates", [])
for candidate in candidates: