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

Add vertex ai image support

Browse source
This commit is contained in:
Sameer Kankute 2025-11-25 13:31:16 +05:30
parent 90850bf6d5
commit 29ab291cf5
5 changed files with 550 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
litellm/images/main.py
View file

@ -19,6 +19,8 @@ from litellm.llms.custom_llm import CustomLLM
#################### Initialize provider clients ####################
llm_http_handler: BaseLLMHTTPHandler = BaseLLMHTTPHandler()
from openai.types.audio.transcription_create_params import FileTypes # type: ignore
from litellm.main import (
azure_chat_completions,
base_llm_aiohttp_handler,
@ -26,7 +28,6 @@ from litellm.main import (
bedrock_image_generation,
openai_chat_completions,
openai_image_variations,
vertex_image_generation,
)
###########################################
@ -36,7 +37,6 @@ from litellm.types.llms.openai import ImageGenerationRequestQuality
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import (
LITELLM_IMAGE_VARIATION_PROVIDERS,
FileTypes,
LlmProviders,
all_litellm_params,
)
@ -344,6 +344,7 @@ def image_generation( # noqa: PLR0915
litellm.LlmProviders.GEMINI,
litellm.LlmProviders.FAL_AI,
litellm.LlmProviders.RUNWAYML,
litellm.LlmProviders.VERTEX_AI,
):
if image_generation_config is None:
raise ValueError(
@ -430,46 +431,6 @@ def image_generation( # noqa: PLR0915
api_base=api_base,
api_key=api_key,
)
elif custom_llm_provider == "vertex_ai":
vertex_ai_project = (
optional_params.pop("vertex_project", None)
or optional_params.pop("vertex_ai_project", None)
or litellm.vertex_project
or get_secret_str("VERTEXAI_PROJECT")
)
vertex_ai_location = (
optional_params.pop("vertex_location", None)
or optional_params.pop("vertex_ai_location", None)
or litellm.vertex_location
or get_secret_str("VERTEXAI_LOCATION")
)
vertex_credentials = (
optional_params.pop("vertex_credentials", None)
or optional_params.pop("vertex_ai_credentials", None)
or get_secret_str("VERTEXAI_CREDENTIALS")
)
api_base = (
api_base
or litellm.api_base
or get_secret_str("VERTEXAI_API_BASE")
or get_secret_str("VERTEX_API_BASE")
)
model_response = vertex_image_generation.image_generation(
model=model,
prompt=prompt,
timeout=timeout,
logging_obj=litellm_logging_obj,
optional_params=optional_params,
model_response=model_response,
vertex_project=vertex_ai_project,
vertex_location=vertex_ai_location,
vertex_credentials=vertex_credentials,
aimg_generation=aimg_generation,
api_base=api_base,
client=client,
)
elif (
custom_llm_provider in litellm._custom_providers
): # Assume custom LLM provider

43
litellm/llms/vertex_ai/image_generation/__init__.py Normal file
View file

@ -0,0 +1,43 @@
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.llms.vertex_ai.common_utils import (
VertexAIModelRoute,
get_vertex_ai_model_route,
)
from .vertex_gemini_transformation import VertexAIGeminiImageGenerationConfig
from .vertex_imagen_transformation import VertexAIImagenImageGenerationConfig
__all__ = [
"VertexAIGeminiImageGenerationConfig",
"VertexAIImagenImageGenerationConfig",
"get_vertex_ai_image_generation_config",
]
def get_vertex_ai_image_generation_config(model: str) -> BaseImageGenerationConfig:
"""
Get the appropriate image generation config for a Vertex AI model.
Routes to the correct transformation class based on the model type:
- Gemini image generation models use generateContent API (VertexAIGeminiImageGenerationConfig)
- Imagen models use predict API (VertexAIImagenImageGenerationConfig)
Args:
model: The model name (e.g., "gemini-2.5-flash-image", "imagegeneration@006")
Returns:
BaseImageGenerationConfig: The appropriate configuration class
"""
# Determine the model route
model_route = get_vertex_ai_model_route(model)
if model_route == VertexAIModelRoute.GEMINI:
# Gemini models use generateContent API
return VertexAIGeminiImageGenerationConfig()
else:
# Default to Imagen for other models (imagegeneration, etc.)
# This includes NON_GEMINI models like imagegeneration@006
return VertexAIImagenImageGenerationConfig()

266
litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py Normal file
View file

@ -0,0 +1,266 @@
import json
import os
from typing import TYPE_CHECKING, Any, Dict, List, Optional
import httpx
import litellm
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIImageGenerationOptionalParams,
)
from litellm.types.utils import ImageObject, ImageResponse
from litellm.utils import ModelResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
"""
Vertex AI Gemini Image Generation Configuration
Uses generateContent API for Gemini image generation models on Vertex AI
Supports models like gemini-2.5-flash-image, gemini-3-pro-image-preview, etc.
"""
def __init__(self) -> None:
BaseImageGenerationConfig.__init__(self)
VertexLLM.__init__(self)
def get_supported_openai_params(
self, model: str
) -> List[OpenAIImageGenerationOptionalParams]:
"""
Gemini image generation supported parameters
"""
return [
"n",
"size",
]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
supported_params = self.get_supported_openai_params(model)
mapped_params = {}
for k, v in non_default_params.items():
if k not in optional_params.keys():
if k in supported_params:
# Map OpenAI parameters to Gemini format
if k == "n":
mapped_params["candidate_count"] = v
elif k == "size":
# Map OpenAI size format to Gemini aspectRatio
mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(v)
else:
mapped_params[k] = v
return mapped_params
def _map_size_to_aspect_ratio(self, size: str) -> str:
"""
Map OpenAI size format to Gemini aspect ratio format
"""
aspect_ratio_map = {
"1024x1024": "1:1",
"1792x1024": "16:9",
"1024x1792": "9:16",
"1280x896": "4:3",
"896x1280": "3:4"
}
return aspect_ratio_map.get(size, "1:1")
def _resolve_vertex_project(self) -> Optional[str]:
return (
getattr(self, "_vertex_project", None)
or os.environ.get("VERTEXAI_PROJECT")
or getattr(litellm, "vertex_project", None)
or get_secret_str("VERTEXAI_PROJECT")
)
def _resolve_vertex_location(self) -> Optional[str]:
return (
getattr(self, "_vertex_location", None)
or os.environ.get("VERTEXAI_LOCATION")
or os.environ.get("VERTEX_LOCATION")
or getattr(litellm, "vertex_location", None)
or get_secret_str("VERTEXAI_LOCATION")
or get_secret_str("VERTEX_LOCATION")
)
def _resolve_vertex_credentials(self) -> Optional[str]:
return (
getattr(self, "_vertex_credentials", None)
or os.environ.get("VERTEXAI_CREDENTIALS")
or getattr(litellm, "vertex_credentials", None)
or os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
or get_secret_str("VERTEXAI_CREDENTIALS")
)
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
"""
Get the complete URL for Vertex AI Gemini generateContent API
"""
vertex_project = self._resolve_vertex_project()
vertex_location = self._resolve_vertex_location()
if not vertex_project or not vertex_location:
raise ValueError("vertex_project and vertex_location are required for Vertex AI")
# Use the model name as provided, handling vertex_ai prefix
model_name = model
if model.startswith("vertex_ai/"):
model_name = model.replace("vertex_ai/", "")
if api_base:
base_url = api_base.rstrip("/")
else:
base_url = f"https://{vertex_location}-aiplatform.googleapis.com"
return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}:generateContent"
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
headers = headers or {}
vertex_project = self._resolve_vertex_project()
vertex_credentials = self._resolve_vertex_credentials()
access_token, _ = self._ensure_access_token(
credentials=vertex_credentials,
project_id=vertex_project,
custom_llm_provider="vertex_ai",
)
return self.set_headers(access_token, headers)
def transform_image_generation_request(
self,
model: str,
prompt: str,
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
"""
Transform the image generation request to Gemini format
Uses generateContent API with responseModalities: ["IMAGE"]
"""
# Prepare messages with the prompt
contents = [
{
"role": "user",
"parts": [{"text": prompt}]
}
]
# Prepare generation config
generation_config: Dict[str, Any] = {
"responseModalities": ["IMAGE"]
}
# Handle image-specific config parameters
image_config: Dict[str, Any] = {}
# Map aspectRatio
if "aspectRatio" in optional_params:
image_config["aspectRatio"] = optional_params["aspectRatio"]
elif "aspect_ratio" in optional_params:
image_config["aspectRatio"] = optional_params["aspect_ratio"]
# Map imageSize (for Gemini 3 Pro)
if "imageSize" in optional_params:
image_config["imageSize"] = optional_params["imageSize"]
elif "image_size" in optional_params:
image_config["imageSize"] = optional_params["image_size"]
if image_config:
generation_config["imageConfig"] = image_config
# Handle candidate_count (n parameter)
if "candidate_count" in optional_params:
generation_config["candidateCount"] = optional_params["candidate_count"]
elif "n" in optional_params:
generation_config["candidateCount"] = optional_params["n"]
request_body: Dict[str, Any] = {
"contents": contents,
"generationConfig": generation_config
}
return request_body
def transform_image_generation_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict,
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ImageResponse:
"""
Transform Gemini image generation response to litellm ImageResponse format
"""
try:
response_data = raw_response.json()
except Exception as e:
raise self.get_error_class(
error_message=f"Error transforming image generation response: {e}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
if not model_response.data:
model_response.data = []
# Gemini image generation models return 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,
))
return model_response

231
litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py Normal file
View file

@ -0,0 +1,231 @@
import json
import os
from typing import TYPE_CHECKING, Any, Dict, List, Optional
import httpx
import litellm
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIImageGenerationOptionalParams,
)
from litellm.types.utils import ImageObject, ImageResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
"""
Vertex AI Imagen Image Generation Configuration
Uses predict API for Imagen models on Vertex AI
Supports models like imagegeneration@006
"""
def __init__(self) -> None:
BaseImageGenerationConfig.__init__(self)
VertexLLM.__init__(self)
def get_supported_openai_params(
self, model: str
) -> List[OpenAIImageGenerationOptionalParams]:
"""
Imagen API supported parameters
"""
return [
"n",
"size"
]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
supported_params = self.get_supported_openai_params(model)
mapped_params = {}
for k, v in non_default_params.items():
if k not in optional_params.keys():
if k in supported_params:
# Map OpenAI parameters to Imagen format
if k == "n":
mapped_params["sampleCount"] = v
elif k == "size":
# Map OpenAI size format to Imagen aspectRatio
mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(v)
else:
mapped_params[k] = v
return mapped_params
def _map_size_to_aspect_ratio(self, size: str) -> str:
"""
Map OpenAI size format to Imagen aspect ratio format
"""
aspect_ratio_map = {
"1024x1024": "1:1",
"1792x1024": "16:9",
"1024x1792": "9:16",
"1280x896": "4:3",
"896x1280": "3:4"
}
return aspect_ratio_map.get(size, "1:1")
def _resolve_vertex_project(self) -> Optional[str]:
return (
getattr(self, "_vertex_project", None)
or os.environ.get("VERTEXAI_PROJECT")
or getattr(litellm, "vertex_project", None)
or get_secret_str("VERTEXAI_PROJECT")
)
def _resolve_vertex_location(self) -> Optional[str]:
return (
getattr(self, "_vertex_location", None)
or os.environ.get("VERTEXAI_LOCATION")
or os.environ.get("VERTEX_LOCATION")
or getattr(litellm, "vertex_location", None)
or get_secret_str("VERTEXAI_LOCATION")
or get_secret_str("VERTEX_LOCATION")
)
def _resolve_vertex_credentials(self) -> Optional[str]:
return (
getattr(self, "_vertex_credentials", None)
or os.environ.get("VERTEXAI_CREDENTIALS")
or getattr(litellm, "vertex_credentials", None)
or os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
or get_secret_str("VERTEXAI_CREDENTIALS")
)
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
"""
Get the complete URL for Vertex AI Imagen predict API
"""
vertex_project = self._resolve_vertex_project()
vertex_location = self._resolve_vertex_location()
if not vertex_project or not vertex_location:
raise ValueError("vertex_project and vertex_location are required for Vertex AI")
# Use the model name as provided, handling vertex_ai prefix
model_name = model
if model.startswith("vertex_ai/"):
model_name = model.replace("vertex_ai/", "")
if api_base:
base_url = api_base.rstrip("/")
else:
base_url = f"https://{vertex_location}-aiplatform.googleapis.com"
return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}:predict"
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
headers = headers or {}
vertex_project = self._resolve_vertex_project()
vertex_credentials = self._resolve_vertex_credentials()
access_token, _ = self._ensure_access_token(
credentials=vertex_credentials,
project_id=vertex_project,
custom_llm_provider="vertex_ai",
)
return self.set_headers(access_token, headers)
def transform_image_generation_request(
self,
model: str,
prompt: str,
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
"""
Transform the image generation request to Imagen format
Uses predict API with instances and parameters
"""
# Default parameters
default_params = {
"sampleCount": 1,
}
# Merge with optional params
parameters = {**default_params, **optional_params}
request_body = {
"instances": [{"prompt": prompt}],
"parameters": parameters,
}
return request_body
def transform_image_generation_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ImageResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict,
optional_params: dict,
litellm_params: dict,
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> ImageResponse:
"""
Transform Imagen image generation response to litellm ImageResponse format
"""
try:
response_data = raw_response.json()
except Exception as e:
raise self.get_error_class(
error_message=f"Error transforming image generation response: {e}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
if not model_response.data:
model_response.data = []
# Imagen format - predictions with generated images
predictions = response_data.get("predictions", [])
for prediction in predictions:
# Imagen returns images as bytesBase64Encoded
if "bytesBase64Encoded" in prediction:
model_response.data.append(ImageObject(
b64_json=prediction["bytesBase64Encoded"],
url=None,
))
return model_response

8
litellm/utils.py
View file

@ -6866,7 +6866,7 @@ def convert_to_dict(message: Union[BaseModel, dict]) -> dict:
dict: The converted message.
"""
if isinstance(message, BaseModel):
return message.model_dump(exclude_none=True)
return message.model_dump(exclude_none=True) # type: ignore
elif isinstance(message, dict):
return message
else:
@ -7671,6 +7671,12 @@ class ProviderConfigManager:
)
return get_runwayml_image_generation_config(model)
elif LlmProviders.VERTEX_AI == provider:
from litellm.llms.vertex_ai.image_generation import (
get_vertex_ai_image_generation_config,
)
return get_vertex_ai_image_generation_config(model)
return None
@staticmethod