import base64 import json import os from io import BufferedRandom, BufferedReader, BytesIO from pathlib import Path from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig 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.images.main import ImageEditOptionalRequestParams from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import FileTypes, ImageObject, ImageResponse, OpenAIImage if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj else: LiteLLMLoggingObj = Any class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): """ Vertex AI Imagen Image Edit Configuration Uses predict API for Imagen models on Vertex AI """ SUPPORTED_PARAMS: list[str] = ["n", "size", "mask"] def __init__(self) -> None: BaseImageEditConfig.__init__(self) VertexLLM.__init__(self) def get_supported_openai_params(self, model: str) -> list[str]: return list(self.SUPPORTED_PARAMS) def map_openai_params( self, image_edit_optional_params: ImageEditOptionalRequestParams, model: str, drop_params: bool, ) -> dict[str, Any]: supported_params: Final = self.get_supported_openai_params(model) filtered_params = {key: value for key, value in image_edit_optional_params.items() if key in supported_params} mapped_params: Final[dict[str, Any]] = {} # Map OpenAI parameters to Imagen format if "n" in filtered_params: mapped_params["sampleCount"] = filtered_params["n"] if "size" in filtered_params: mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(filtered_params["size"]) if "mask" in filtered_params: mapped_params["mask"] = filtered_params["mask"] return mapped_params def _resolve_vertex_project(self) -> str | None: 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) -> str | None: 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) -> str | None: 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 validate_environment( self, headers: dict, model: str, api_key: str | None = None, litellm_params: dict | None = None, api_base: str | None = None, ) -> dict: headers = headers or {} litellm_params = litellm_params or {} _api_base: Final = litellm_params.get("api_base") or api_base if _api_base is not None: return headers vertex_project: Final = 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 get_complete_url( self, model: str, api_base: str | None, litellm_params: dict, ) -> str: """ Get the complete URL for Vertex AI Imagen predict API """ vertex_project: Final = self.safe_get_vertex_ai_project(litellm_params) or self._resolve_vertex_project() vertex_location: Final = 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") # 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 = get_vertex_base_url(vertex_location) return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}:predict" def transform_image_edit_request( self, model: str, prompt: str | None, image: FileTypes | None, image_edit_optional_request_params: dict[str, Any], litellm_params: GenericLiteLLMParams, headers: dict, ) -> tuple[dict[str, Any], RequestFiles | None]: # Prepare reference images in the correct Imagen format if image is None: raise ValueError("Vertex AI Imagen image edit requires at least one reference image.") reference_images: Final = self._prepare_reference_images(image, image_edit_optional_request_params) if not reference_images: raise ValueError("Vertex AI Imagen image edit requires at least one reference image.") if prompt is None: raise ValueError("Vertex AI Imagen image edit requires a prompt.") # Correct Imagen instances format instances: Final = [{"prompt": prompt, "referenceImages": reference_images}] # Extract OpenAI parameters and set sensible defaults for Vertex AI-specific parameters sample_count: Final = image_edit_optional_request_params.get("sampleCount", 1) # Use sensible defaults for Vertex AI-specific parameters (not exposed to users) edit_mode: Final = "EDIT_MODE_INPAINT_INSERTION" # Default edit mode base_steps: Final = 50 # Default number of steps # Imagen parameters with correct structure parameters: Final = { "sampleCount": sample_count, "editMode": edit_mode, "editConfig": {"baseSteps": base_steps}, } # Set default values for Vertex AI-specific parameters (not configurable by users via OpenAI API) parameters["guidanceScale"] = 7.5 # Default guidance scale parameters["seed"] = None # Let Vertex AI choose random seed request_body: Final[dict[str, Any]] = { "instances": instances, "parameters": parameters, } payload: Final[Any] = json.dumps(request_body) empty_files: Final = cast(RequestFiles, []) return cast(tuple[dict[str, Any], RequestFiles | None], (payload, empty_files)) def transform_image_edit_response( self, model: str, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, ) -> ImageResponse: model_response: Final = ImageResponse() try: response_json: Final = raw_response.json() except Exception as exc: raise self.get_error_class( error_message=f"Error transforming image edit response: {exc}", status_code=raw_response.status_code, headers=raw_response.headers, ) predictions: Final = response_json.get("predictions", []) data_list: Final[list[ImageObject]] = [] for prediction in predictions: # Imagen returns images as bytesBase64Encoded if "bytesBase64Encoded" in prediction: data_list.append( ImageObject( b64_json=prediction["bytesBase64Encoded"], url=None, ) ) model_response.data = cast(list[OpenAIImage], data_list) return model_response def _map_size_to_aspect_ratio(self, size: str) -> str: """Map OpenAI size format to Imagen aspect ratio format""" aspect_ratio_map: Final = { "1024x1024": "1:1", "1792x1024": "16:9", "1024x1792": "9:16", "1280x896": "4:3", "896x1280": "3:4", } return aspect_ratio_map.get(size, "1:1") def _prepare_reference_images( self, image: FileTypes | list[FileTypes], image_edit_optional_request_params: dict[str, Any], ) -> list[dict[str, Any]]: """ Prepare reference images in the correct Imagen API format """ images: list[FileTypes] if isinstance(image, list): images = image else: images = [image] reference_images: Final[list[dict[str, Any]]] = [] for idx, img in enumerate(images): if img is None: continue image_bytes = self._read_all_bytes(img) base64_data = base64.b64encode(image_bytes).decode("utf-8") # Create reference image structure reference_image = { "referenceType": "REFERENCE_TYPE_RAW", "referenceId": idx + 1, "referenceImage": {"bytesBase64Encoded": base64_data}, } reference_images.append(reference_image) # Handle mask image if provided (for inpainting) mask_image: Final = image_edit_optional_request_params.get("mask") if mask_image is not None: mask_bytes: Final = self._read_all_bytes(mask_image) mask_base64: Final = base64.b64encode(mask_bytes).decode("utf-8") mask_reference: Final = { "referenceType": "REFERENCE_TYPE_MASK", "referenceId": len(reference_images) + 1, "referenceImage": {"bytesBase64Encoded": mask_base64}, "maskImageConfig": { "maskMode": "MASK_MODE_USER_PROVIDED", "dilation": 0.03, # Default dilation value (not configurable via OpenAI API) }, } reference_images.append(mask_reference) return reference_images def _read_all_bytes(self, image: Any, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> bytes: if depth > max_depth: raise ValueError( f"Max recursion depth {max_depth} reached while reading image bytes for Vertex AI Imagen image edit." ) if isinstance(image, (list, tuple)): for item in image: if item is not None: return self._read_all_bytes(item, depth=depth + 1, max_depth=max_depth) raise ValueError("Unsupported image type for Vertex AI Imagen image edit.") if isinstance(image, dict): for key in ("data", "bytes", "content"): if key in image and image[key] is not None: value = image[key] if isinstance(value, str): try: return base64.b64decode(value) except Exception: continue return self._read_all_bytes(value, depth=depth + 1, max_depth=max_depth) if "path" in image: return self._read_all_bytes(image["path"], depth=depth + 1, max_depth=max_depth) if isinstance(image, bytes): return image if isinstance(image, bytearray): return bytes(image) if isinstance(image, BytesIO): current_pos: Final = image.tell() image.seek(0) data = image.read() image.seek(current_pos) return data if isinstance(image, (BufferedReader, BufferedRandom)): stream_pos: int | None = None try: stream_pos = image.tell() except Exception: stream_pos = None if stream_pos is not None: image.seek(0) data = image.read() if stream_pos is not None: image.seek(stream_pos) return data if isinstance(image, str): raise ValueError( "Unsupported image input: plain string values are not accepted for " "Vertex AI Imagen image edit. Provide image bytes or a file-like object." ) if isinstance(image, Path): raise ValueError( "Unsupported image input: filesystem paths are not accepted for " "Vertex AI Imagen image edit. Provide image bytes or a file-like object." ) if hasattr(image, "read"): data = image.read() if isinstance(data, str): data = data.encode("utf-8") return data raise ValueError(f"Unsupported image type for Vertex AI Imagen image edit. Got type={type(image)}")