From 29ab291cf5c6e4895e587b982b26792d730da34b Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 25 Nov 2025 13:31:16 +0530 Subject: [PATCH 1/4] Add vertex ai image support --- litellm/images/main.py | 45 +-- .../vertex_ai/image_generation/__init__.py | 43 +++ .../vertex_gemini_transformation.py | 266 ++++++++++++++++++ .../vertex_imagen_transformation.py | 231 +++++++++++++++ litellm/utils.py | 8 +- 5 files changed, 550 insertions(+), 43 deletions(-) create mode 100644 litellm/llms/vertex_ai/image_generation/__init__.py create mode 100644 litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py create mode 100644 litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py diff --git a/litellm/images/main.py b/litellm/images/main.py index 333a751b04..786136e669 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -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 diff --git a/litellm/llms/vertex_ai/image_generation/__init__.py b/litellm/llms/vertex_ai/image_generation/__init__.py new file mode 100644 index 0000000000..a6f6156167 --- /dev/null +++ b/litellm/llms/vertex_ai/image_generation/__init__.py @@ -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() + diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py new file mode 100644 index 0000000000..0a87dea997 --- /dev/null +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -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 + diff --git a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py new file mode 100644 index 0000000000..275c547c8d --- /dev/null +++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py @@ -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 + diff --git a/litellm/utils.py b/litellm/utils.py index 1ec2576d35..3bc570bd70 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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 From f52f05748dc688895c7a26a4b65d8ea8d8c0b59c Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 25 Nov 2025 13:31:50 +0530 Subject: [PATCH 2/4] Update docs related to vertex ai image gen --- .../my-website/docs/providers/vertex_image.md | 73 +++++++++++++++++-- 1 file changed, 66 insertions(+), 7 deletions(-) diff --git a/docs/my-website/docs/providers/vertex_image.md b/docs/my-website/docs/providers/vertex_image.md index 27e584cb22..c4d5d55408 100644 --- a/docs/my-website/docs/providers/vertex_image.md +++ b/docs/my-website/docs/providers/vertex_image.md @@ -1,18 +1,65 @@ # Vertex AI Image Generation -Vertex AI Image Generation uses Google's Imagen models to generate high-quality images from text descriptions. +Vertex AI supports two types of image generation: + +1. **Gemini Image Generation Models** (Nano Banana 🍌) - Conversational image generation using `generateContent` API +2. **Imagen Models** - Traditional image generation using `predict` API | Property | Details | |----------|---------| -| Description | Vertex AI Image Generation uses Google's Imagen models to generate high-quality images from text descriptions. | +| Description | Vertex AI Image Generation supports both Gemini image generation models | | Provider Route on LiteLLM | `vertex_ai/` | | Provider Doc | [Google Cloud Vertex AI Image Generation ↗](https://cloud.google.com/vertex-ai/docs/generative-ai/image/generate-images) | +| Gemini Image Generation Docs | [Gemini Image Generation ↗](https://ai.google.dev/gemini-api/docs/image-generation) | ## Quick Start -### LiteLLM Python SDK +### Gemini Image Generation Models -```python showLineNumbers title="Basic Image Generation" +Gemini image generation models support conversational image creation with features like: +- Text-to-Image generation +- Image editing (text + image → image) +- Multi-turn image refinement +- High-fidelity text rendering +- Up to 4K resolution (Gemini 3 Pro) + +```python showLineNumbers title="Gemini 2.5 Flash Image" +import litellm + +# Generate a single image +response = await litellm.aimage_generation( + prompt="A nano banana dish in a fancy restaurant with a Gemini theme", + model="vertex_ai/gemini-2.5-flash-image", + vertex_ai_project="your-project-id", + vertex_ai_location="us-central1", + n=1, + size="1024x1024", +) + +print(response.data[0].b64_json) # Gemini returns base64 images +``` + +```python showLineNumbers title="Gemini 3 Pro Image Preview (4K output)" +import litellm + +# Generate high-resolution image +response = await litellm.aimage_generation( + prompt="Da Vinci style anatomical sketch of a dissected Monarch butterfly", + model="vertex_ai/gemini-3-pro-image-preview", + vertex_ai_project="your-project-id", + vertex_ai_location="us-central1", + n=1, + size="1024x1024", + # Optional: specify image size for Gemini 3 Pro + # imageSize="4K", # Options: "1K", "2K", "4K" +) + +print(response.data[0].b64_json) +``` + +### Imagen Models + +```python showLineNumbers title="Imagen Image Generation" import litellm # Generate a single image @@ -21,9 +68,11 @@ response = await litellm.aimage_generation( model="vertex_ai/imagen-4.0-generate-001", vertex_ai_project="your-project-id", vertex_ai_location="us-central1", + n=1, + size="1024x1024", ) -print(response.data[0].url) +print(response.data[0].b64_json) # Imagen also returns base64 images ``` ### LiteLLM Proxy @@ -70,6 +119,18 @@ print(response.data[0].url) ## Supported Models +### Gemini Image Generation Models + +- `vertex_ai/gemini-2.5-flash-image` - Fast, efficient image generation (1024px resolution) +- `vertex_ai/gemini-3-pro-image-preview` - Advanced model with 4K output, Google Search grounding, and thinking mode +- `vertex_ai/gemini-2.0-flash-preview-image` - Preview model +- `vertex_ai/gemini-2.5-flash-image-preview` - Preview model + +### Imagen Models + +- `vertex_ai/imagegeneration@006` - Legacy Imagen model +- `vertex_ai/imagen-4.0-generate-001` - Latest Imagen model +- `vertex_ai/imagen-3.0-generate-001` - Imagen 3.0 model :::tip @@ -77,7 +138,5 @@ print(response.data[0].url) ::: -LiteLLM supports all Vertex AI Imagen models available through Google Cloud. - For the complete and up-to-date list of supported models, visit: [https://models.litellm.ai/](https://models.litellm.ai/) From 883cfaeeafa2ecb16f91a540be63dfbfa644001c Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 25 Nov 2025 13:32:13 +0530 Subject: [PATCH 3/4] Add tests --- .../image_gen_tests/test_image_generation.py | 32 ++ ...rtex_ai_image_generation_transformation.py | 457 ++++++++++++++++++ 2 files changed, 489 insertions(+) create mode 100644 tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index 24d5843293..add60c755b 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -119,6 +119,38 @@ class TestVertexImageGeneration(BaseImageGenTest): } +class TestVertexAIGeminiImageGeneration(BaseImageGenTest): + """Test Gemini image generation models (Nano Banana)""" + def get_base_image_generation_call_args(self) -> dict: + # comment this when running locally + load_vertex_ai_credentials() + + litellm.in_memory_llm_clients_cache = InMemoryCache() + return { + "model": "vertex_ai/gemini-2.5-flash-image", + "vertex_ai_project": "pathrise-convert-1606954137718", + "vertex_ai_location": "us-central1", + "n": 1, + "size": "1024x1024", + } + + +class TestVertexAIGemini3ProImageGeneration(BaseImageGenTest): + """Test Gemini 3 Pro image generation model""" + def get_base_image_generation_call_args(self) -> dict: + # comment this when running locally + load_vertex_ai_credentials() + + litellm.in_memory_llm_clients_cache = InMemoryCache() + return { + "model": "vertex_ai/gemini-3-pro-image-preview", + "vertex_ai_project": "pathrise-convert-1606954137718", + "vertex_ai_location": "us-central1", + "n": 1, + "size": "1024x1024", + } + + class TestBedrockNovaCanvasTextToImage(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: litellm.in_memory_llm_clients_cache = InMemoryCache() diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py new file mode 100644 index 0000000000..9f33400594 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -0,0 +1,457 @@ +import os +import sys +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.llms.vertex_ai.image_generation import ( + get_vertex_ai_image_generation_config, +) +from litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation import ( + VertexAIGeminiImageGenerationConfig, +) +from litellm.llms.vertex_ai.image_generation.vertex_imagen_transformation import ( + VertexAIImagenImageGenerationConfig, +) + + +class TestVertexAIGeminiImageGenerationConfig: + def setup_method(self): + """Set up test fixtures""" + self.config = VertexAIGeminiImageGenerationConfig() + + def test_get_supported_openai_params(self): + """Test get_supported_openai_params returns correct params""" + supported = self.config.get_supported_openai_params("gemini-2.5-flash-image") + assert "n" in supported + assert "size" in supported + + def test_map_openai_params_n(self): + """Test mapping n parameter to candidate_count""" + non_default_params = {"n": 3} + optional_params = {} + result = self.config.map_openai_params( + non_default_params, optional_params, "gemini-2.5-flash-image", False + ) + assert result.get("candidate_count") == 3 + + def test_map_openai_params_size(self): + """Test mapping size parameter to aspectRatio""" + non_default_params = {"size": "1024x1024"} + optional_params = {} + result = self.config.map_openai_params( + non_default_params, optional_params, "gemini-2.5-flash-image", False + ) + assert result.get("aspectRatio") == "1:1" + + def test_map_openai_params_size_16_9(self): + """Test mapping 16:9 size""" + non_default_params = {"size": "1792x1024"} + optional_params = {} + result = self.config.map_openai_params( + non_default_params, optional_params, "gemini-2.5-flash-image", False + ) + assert result.get("aspectRatio") == "16:9" + + def test_map_size_to_aspect_ratio(self): + """Test size to aspect ratio mapping""" + assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" + assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" + assert self.config._map_size_to_aspect_ratio("1024x1792") == "9:16" + assert self.config._map_size_to_aspect_ratio("1280x896") == "4:3" + assert self.config._map_size_to_aspect_ratio("896x1280") == "3:4" + assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default + + def test_transform_image_generation_request_basic(self): + """Test basic request transformation""" + request = self.config.transform_image_generation_request( + model="gemini-2.5-flash-image", + prompt="A nano banana", + optional_params={}, + litellm_params={}, + headers={}, + ) + assert "contents" in request + assert "generationConfig" in request + assert request["generationConfig"]["responseModalities"] == ["IMAGE"] + assert request["contents"][0]["parts"][0]["text"] == "A nano banana" + + def test_transform_image_generation_request_with_aspect_ratio(self): + """Test request transformation with aspectRatio""" + request = self.config.transform_image_generation_request( + model="gemini-2.5-flash-image", + prompt="A nano banana", + optional_params={"aspectRatio": "16:9"}, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9" + + def test_transform_image_generation_request_with_image_size(self): + """Test request transformation with imageSize (Gemini 3 Pro)""" + request = self.config.transform_image_generation_request( + model="gemini-3-pro-image-preview", + prompt="A nano banana", + optional_params={"imageSize": "4K"}, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["imageConfig"]["imageSize"] == "4K" + + def test_transform_image_generation_request_with_candidate_count(self): + """Test request transformation with candidate_count""" + request = self.config.transform_image_generation_request( + model="gemini-2.5-flash-image", + prompt="A nano banana", + optional_params={"candidate_count": 2}, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["candidateCount"] == 2 + + def test_transform_image_generation_request_with_n(self): + """Test request transformation with n parameter""" + request = self.config.transform_image_generation_request( + model="gemini-2.5-flash-image", + prompt="A nano banana", + optional_params={"n": 2}, + litellm_params={}, + headers={}, + ) + assert request["generationConfig"]["candidateCount"] == 2 + + def test_transform_image_generation_response(self): + """Test response transformation""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "base64_encoded_image_data", + } + } + ] + } + } + ] + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="gemini-2.5-flash-image", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 1 + assert result.data[0].b64_json == "base64_encoded_image_data" + assert result.data[0].url is None + + def test_transform_image_generation_response_multiple_images(self): + """Test response transformation with multiple images""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "image1", + } + }, + { + "inlineData": { + "mimeType": "image/png", + "data": "image2", + } + }, + ] + } + } + ] + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="gemini-2.5-flash-image", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 2 + assert result.data[0].b64_json == "image1" + assert result.data[1].b64_json == "image2" + + +class TestVertexAIImagenImageGenerationConfig: + def setup_method(self): + """Set up test fixtures""" + self.config = VertexAIImagenImageGenerationConfig() + + def test_get_supported_openai_params(self): + """Test get_supported_openai_params returns correct params""" + supported = self.config.get_supported_openai_params("imagegeneration@006") + assert "n" in supported + assert "size" in supported + + def test_map_openai_params_n(self): + """Test mapping n parameter to sampleCount""" + non_default_params = {"n": 3} + optional_params = {} + result = self.config.map_openai_params( + non_default_params, optional_params, "imagegeneration@006", False + ) + assert result.get("sampleCount") == 3 + + def test_map_openai_params_size(self): + """Test mapping size parameter to aspectRatio""" + non_default_params = {"size": "1024x1024"} + optional_params = {} + result = self.config.map_openai_params( + non_default_params, optional_params, "imagegeneration@006", False + ) + assert result.get("aspectRatio") == "1:1" + + def test_map_size_to_aspect_ratio(self): + """Test size to aspect ratio mapping""" + assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1" + assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9" + assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default + + def test_transform_image_generation_request_basic(self): + """Test basic request transformation""" + request = self.config.transform_image_generation_request( + model="imagegeneration@006", + prompt="A cat", + optional_params={}, + litellm_params={}, + headers={}, + ) + assert "instances" in request + assert "parameters" in request + assert request["instances"][0]["prompt"] == "A cat" + assert request["parameters"]["sampleCount"] == 1 + + def test_transform_image_generation_request_with_params(self): + """Test request transformation with parameters""" + request = self.config.transform_image_generation_request( + model="imagegeneration@006", + prompt="A cat", + optional_params={"sampleCount": 2, "aspectRatio": "16:9"}, + litellm_params={}, + headers={}, + ) + assert request["parameters"]["sampleCount"] == 2 + assert request["parameters"]["aspectRatio"] == "16:9" + + def test_transform_image_generation_response(self): + """Test response transformation""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "predictions": [ + {"bytesBase64Encoded": "base64_encoded_image_data"} + ] + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="imagegeneration@006", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 1 + assert result.data[0].b64_json == "base64_encoded_image_data" + assert result.data[0].url is None + + def test_transform_image_generation_response_multiple_images(self): + """Test response transformation with multiple images""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "predictions": [ + {"bytesBase64Encoded": "image1"}, + {"bytesBase64Encoded": "image2"}, + ] + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="imagegeneration@006", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 2 + assert result.data[0].b64_json == "image1" + assert result.data[1].b64_json == "image2" + + +class TestGetVertexAIImageGenerationConfig: + """Test the router function that selects the correct config""" + + def test_get_gemini_model_config(self): + """Test that Gemini models return Gemini config""" + config = get_vertex_ai_image_generation_config("gemini-2.5-flash-image") + assert isinstance(config, VertexAIGeminiImageGenerationConfig) + + config = get_vertex_ai_image_generation_config("gemini-3-pro-image-preview") + assert isinstance(config, VertexAIGeminiImageGenerationConfig) + + config = get_vertex_ai_image_generation_config( + "vertex_ai/gemini-2.5-flash-image" + ) + assert isinstance(config, VertexAIGeminiImageGenerationConfig) + + def test_get_imagen_model_config(self): + """Test that Imagen models return Imagen config""" + config = get_vertex_ai_image_generation_config("imagegeneration@006") + assert isinstance(config, VertexAIImagenImageGenerationConfig) + + config = get_vertex_ai_image_generation_config("imagen-4.0-generate-001") + assert isinstance(config, VertexAIImagenImageGenerationConfig) + + config = get_vertex_ai_image_generation_config( + "vertex_ai/imagegeneration@006" + ) + assert isinstance(config, VertexAIImagenImageGenerationConfig) + + def test_get_non_gemini_model_config(self): + """Test that non-Gemini models default to Imagen config""" + config = get_vertex_ai_image_generation_config("some-other-model") + assert isinstance(config, VertexAIImagenImageGenerationConfig) + + +class TestVertexAIImageGenerationIntegration: + """Integration tests for Vertex AI image generation""" + + @pytest.mark.skipif( + not os.getenv("VERTEXAI_PROJECT"), + reason="Vertex AI credentials not set", + ) + def test_gemini_image_generation_config_validation(self): + """Test that Gemini config can validate environment""" + config = VertexAIGeminiImageGenerationConfig() + with patch.object( + config, "_resolve_vertex_project", return_value="test-project" + ), patch.object( + config, "_resolve_vertex_location", return_value="us-central1" + ), patch.object( + config, "_ensure_access_token", return_value=("token", None) + ): + headers = config.validate_environment( + headers={}, + model="gemini-2.5-flash-image", + messages=[], + optional_params={}, + litellm_params={}, + ) + assert "Authorization" in headers + + @pytest.mark.skipif( + not os.getenv("VERTEXAI_PROJECT"), + reason="Vertex AI credentials not set", + ) + def test_imagen_image_generation_config_validation(self): + """Test that Imagen config can validate environment""" + config = VertexAIImagenImageGenerationConfig() + with patch.object( + config, "_resolve_vertex_project", return_value="test-project" + ), patch.object( + config, "_resolve_vertex_location", return_value="us-central1" + ), patch.object( + config, "_ensure_access_token", return_value=("token", None) + ): + headers = config.validate_environment( + headers={}, + model="imagegeneration@006", + messages=[], + optional_params={}, + litellm_params={}, + ) + assert "Authorization" in headers + + def test_gemini_get_complete_url(self): + """Test Gemini config URL generation""" + config = VertexAIGeminiImageGenerationConfig() + with patch.object( + config, "_resolve_vertex_project", return_value="test-project" + ), patch.object( + config, "_resolve_vertex_location", return_value="us-central1" + ): + url = config.get_complete_url( + api_base=None, + api_key=None, + model="gemini-2.5-flash-image", + optional_params={}, + litellm_params={}, + ) + assert "test-project" in url + assert "us-central1" in url + assert "gemini-2.5-flash-image" in url + assert "generateContent" in url + + def test_imagen_get_complete_url(self): + """Test Imagen config URL generation""" + config = VertexAIImagenImageGenerationConfig() + with patch.object( + config, "_resolve_vertex_project", return_value="test-project" + ), patch.object( + config, "_resolve_vertex_location", return_value="us-central1" + ): + url = config.get_complete_url( + api_base=None, + api_key=None, + model="imagegeneration@006", + optional_params={}, + litellm_params={}, + ) + assert "test-project" in url + assert "us-central1" in url + assert "imagegeneration@006" in url + assert "predict" in url + From b0d511143c952d523ab5302ac6031618af1d69d9 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 25 Nov 2025 13:36:20 +0530 Subject: [PATCH 4/4] remove unsused imports --- .../vertex_ai/image_generation/vertex_gemini_transformation.py | 2 -- .../vertex_ai/image_generation/vertex_imagen_transformation.py | 3 +-- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 0a87dea997..149e0850bf 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -1,4 +1,3 @@ -import json import os from typing import TYPE_CHECKING, Any, Dict, List, Optional @@ -15,7 +14,6 @@ from litellm.types.llms.openai import ( 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 diff --git a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py index 275c547c8d..8c4ad5dd42 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py @@ -1,6 +1,5 @@ -import json import os -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, List, Optional import httpx