import os from typing import TYPE_CHECKING, Any, List, Optional import httpx import litellm from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) from litellm.llms.vertex_ai.common_utils import get_vertex_base_url 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 """ # 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 a custom api_base is provided, use it directly # This allows users to use proxies or mock endpoints if api_base: return api_base.rstrip("/") # First check litellm_params (where vertex_ai_project/vertex_ai_location are passed) # then fall back to environment variables and other sources vertex_project = ( self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() ) vertex_location = ( self.safe_get_vertex_ai_location(litellm_params) or self._resolve_vertex_location() ) if not vertex_project or not vertex_location: raise ValueError( "vertex_project and vertex_location are required for Vertex AI" ) base_url = get_vertex_base_url(vertex_location) 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 {} # If a custom api_base is provided, skip credential validation # This allows users to use proxies or mock endpoints without needing Vertex AI credentials _api_base = litellm_params.get("api_base") or api_base if _api_base is not None: return headers # First check litellm_params (where vertex_ai_project/vertex_ai_credentials are passed) # then fall back to environment variables and other sources vertex_project = ( self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() ) vertex_credentials = ( self.safe_get_vertex_ai_credentials(litellm_params) or 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