From b60e30a58b13112ab506a3718cc5d9c528f79a86 Mon Sep 17 00:00:00 2001 From: Albert Sebastian Date: Wed, 8 Apr 2026 15:27:34 +0530 Subject: [PATCH] feat(vertex_ai): support image_size (2K/4K) for Gemini image generation Fixes #24621 The Gemini API supports `imageConfig.imageSize` to control output resolution (e.g., "2K", "4K"), but LiteLLM had no way to pass this through. The `extra_body` approach doesn't work because `generationConfig` is rebuilt from scratch in the transformation layer. Changes: - Vertex AI Gemini image edit: add `imageSize` to `SUPPORTED_PARAMS` and include it in `generationConfig.image_config.image_size` - Google AI Studio image gen: same support for Gemini models - Both: accept `size` param as either OpenAI format ("1024x1024" -> aspect_ratio) or Gemini format ("2K" -> image_size) - Both: map OpenAI `quality="hd"` to `imageSize="2K"` as a convenient alternative - Add tests for imageSize, combined aspect_ratio+imageSize, size="2K", and quality="hd" mappings --- .../gemini/image_generation/transformation.py | 28 ++++++-- .../vertex_gemini_transformation.py | 25 +++++-- ...est_vertex_ai_image_edit_transformation.py | 67 +++++++++++++++++++ 3 files changed, 112 insertions(+), 8 deletions(-) diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index b094fc133d7..e28eb5d1797 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -36,7 +36,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): Google AI Imagen API supported parameters https://ai.google.dev/gemini-api/docs/imagen """ - return ["n", "size"] + return ["n", "size", "quality"] def map_openai_params( self, @@ -55,8 +55,17 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): if k == "n": mapped_params["sampleCount"] = v elif k == "size": - # Map OpenAI size format to Google aspectRatio - mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(v) + if isinstance(v, str) and "x" in v: + # OpenAI format like "1024x1024" -> map to aspect ratio + mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(v) + else: + # Gemini image_size format like "1K", "2K", "4K" + mapped_params["imageSize"] = v + elif k == "quality": + # Map OpenAI quality to Gemini imageSize + # "hd" -> "2K", "standard" -> "1K" + if v == "hd" and "imageSize" not in mapped_params: + mapped_params["imageSize"] = "2K" else: mapped_params[k] = v return mapped_params @@ -180,9 +189,20 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): """ # For Gemini Flash Image Preview models, use standard Gemini format if "gemini" in model: + generation_config: dict = {"response_modalities": ["IMAGE", "TEXT"]} + + # Build image_config from mapped params (aspectRatio, imageSize) + image_config: dict = {} + if "aspectRatio" in optional_params: + image_config["aspect_ratio"] = optional_params["aspectRatio"] + if "imageSize" in optional_params: + image_config["image_size"] = optional_params["imageSize"] + if image_config: + generation_config["image_config"] = image_config + request_body: dict = { "contents": [{"parts": [{"text": prompt}]}], - "generationConfig": {"response_modalities": ["IMAGE", "TEXT"]}, + "generationConfig": generation_config, } return request_body else: diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index de7f234a861..2944dc18f78 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -32,7 +32,7 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): Uses generateContent API for Gemini models on Vertex AI """ - SUPPORTED_PARAMS: List[str] = ["size"] + SUPPORTED_PARAMS: List[str] = ["size", "quality"] def __init__(self) -> None: BaseImageEditConfig.__init__(self) @@ -57,9 +57,22 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): mapped_params: Dict[str, Any] = {} if "size" in filtered_params: - mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio( - filtered_params["size"] # type: ignore[arg-type] - ) + size_value = filtered_params["size"] + if isinstance(size_value, str) and "x" in size_value: + # OpenAI format like "1024x1024" -> map to aspect ratio + mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio( + size_value + ) + else: + # Gemini image_size format like "1K", "2K", "4K" + mapped_params["imageSize"] = size_value + + if "quality" in filtered_params: + # Map OpenAI quality to Gemini imageSize + # "hd" -> "2K", "standard" -> "1K" + quality = filtered_params["quality"] + if quality == "hd" and "imageSize" not in mapped_params: + mapped_params["imageSize"] = "2K" return mapped_params @@ -195,6 +208,10 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): image_config["aspect_ratio"] = image_edit_optional_request_params[ "aspectRatio" ] + if "imageSize" in image_edit_optional_request_params: + image_config["image_size"] = image_edit_optional_request_params[ + "imageSize" + ] if image_config: generation_config["image_config"] = image_config diff --git a/tests/test_litellm/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py b/tests/test_litellm/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py index c231904e710..00fb7740b6b 100644 --- a/tests/test_litellm/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/image_edit/test_vertex_ai_image_edit_transformation.py @@ -126,6 +126,73 @@ class TestVertexAIGeminiImageEditTransformation: "utf-8" ) + def test_transform_image_edit_request_with_image_size(self) -> None: + """Test that imageSize is included in image_config""" + image_bytes = b"fake_image_data" + image = BytesIO(image_bytes) + optional_params = { + "imageSize": "2K", + } + + request_body_str, files = self.config.transform_image_edit_request( + model=self.model, + prompt=self.prompt, + image=image, + image_edit_optional_request_params=optional_params, + litellm_params=MagicMock(), + headers={}, + ) + + request_body = json.loads(request_body_str) + generation_config = request_body["generationConfig"] + assert "image_config" in generation_config + assert generation_config["image_config"]["image_size"] == "2K" + + def test_transform_image_edit_request_with_aspect_ratio_and_image_size(self) -> None: + """Test that both aspectRatio and imageSize are included in image_config""" + image_bytes = b"fake_image_data" + image = BytesIO(image_bytes) + optional_params = { + "aspectRatio": "16:9", + "imageSize": "2K", + } + + request_body_str, files = self.config.transform_image_edit_request( + model=self.model, + prompt=self.prompt, + image=image, + image_edit_optional_request_params=optional_params, + litellm_params=MagicMock(), + headers={}, + ) + + request_body = json.loads(request_body_str) + image_config = request_body["generationConfig"]["image_config"] + assert image_config["aspect_ratio"] == "16:9" + assert image_config["image_size"] == "2K" + + def test_map_openai_params_size_as_resolution(self) -> None: + """Test that size='2K' maps to imageSize instead of aspectRatio""" + optional_params: Dict[str, object] = {"size": "2K"} + mapped = self.config.map_openai_params( + image_edit_optional_params=optional_params, # type: ignore[arg-type] + model=self.model, + drop_params=False, + ) + assert "imageSize" in mapped + assert mapped["imageSize"] == "2K" + assert "aspectRatio" not in mapped + + def test_map_openai_params_quality_hd(self) -> None: + """Test that quality='hd' maps to imageSize='2K'""" + optional_params: Dict[str, object] = {"quality": "hd"} + mapped = self.config.map_openai_params( + image_edit_optional_params=optional_params, # type: ignore[arg-type] + model=self.model, + drop_params=False, + ) + assert mapped["imageSize"] == "2K" + def test_transform_image_edit_request_without_image_raises(self) -> None: """Test that missing image raises ValueError""" optional_params = {}