diff --git a/litellm/images/main.py b/litellm/images/main.py index 333a751b045..786136e6699 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 00000000000..a6f6156167a --- /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 00000000000..0a87dea997a --- /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 00000000000..275c547c8dd --- /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 1ec2576d356..3bc570bd700 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