From 5cfcc98d2ea7c0bf834f60d699e9b4fbf49a5010 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 26 Nov 2025 17:36:37 -0800 Subject: [PATCH] fix img gen --- litellm/images/main.py | 7 ++++- .../vertex_gemini_transformation.py | 27 ++++++++++++------- .../vertex_imagen_transformation.py | 27 ++++++++++++------- 3 files changed, 42 insertions(+), 19 deletions(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index 878fce83f1..eacd477829 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -10,7 +10,8 @@ from litellm import client, exception_type, get_litellm_params from litellm.constants import DEFAULT_IMAGE_ENDPOINT_MODEL from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT from litellm.exceptions import LiteLLMUnknownProvider -from litellm.litellm_core_utils.litellm_logging import Logging, Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.mock_functions import mock_image_generation from litellm.llms.base_llm import BaseImageEditConfig, BaseImageGenerationConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -351,6 +352,10 @@ def image_generation( # noqa: PLR0915 f"image generation config is not supported for {custom_llm_provider}" ) + # Resolve api_base from litellm.api_base if not explicitly provided + _api_base = api_base or litellm.api_base + litellm_params_dict["api_base"] = _api_base + return llm_http_handler.image_generation_handler( api_key=api_key, model=model, 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 c863c2f569..416c611d86 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -122,6 +122,16 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): """ Get the complete URL for Vertex AI Gemini generateContent 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() @@ -130,15 +140,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): 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" + 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" @@ -153,6 +155,13 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): 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() 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 4d86f7ba36..33f416f9ca 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py @@ -122,6 +122,16 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): """ 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() @@ -130,15 +140,7 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): 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" + 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" @@ -153,6 +155,13 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): 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()