mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
fix(vertex-ai): preserve image size aspect ratios
This commit is contained in:
shrey kharbanda
parent
f96200bcea
commit
acbbddd55b
10 changed files with 349 additions and 94 deletions
|
|
@ -18,6 +18,25 @@ from litellm.secret_managers.main import get_secret_str
|
|||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse
|
||||
|
||||
# 2.5 Flash has no configurable imageSize and supports ten aspect ratios, excluding these four
|
||||
# https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/2-5-flash-image
|
||||
_GEMINI_2_5_FLASH_MODEL_FAMILY: Final = "2.5-flash"
|
||||
_GEMINI_2_5_FLASH_UNSUPPORTED_ASPECT_RATIOS: Final = frozenset(("1:4", "4:1", "1:8", "8:1"))
|
||||
|
||||
# 3.1 Flash supports the 512 tier, while other Gemini image models start at 1K
|
||||
# https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/3-1-flash-image
|
||||
_GEMINI_512_IMAGE_SIZE_MODEL: Final = "gemini-3.1-flash-image"
|
||||
|
||||
# Flash Lite supports only 1K, even when the requested dimensions are larger
|
||||
# https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/3-1-flash-lite-image
|
||||
_GEMINI_1K_ONLY_IMAGE_SIZE_MODEL: Final = "gemini-3.1-flash-lite-image"
|
||||
|
||||
# LiteLLM chooses the nearest resolution tier using the side length of a square with the same total number of pixels
|
||||
# as the requested image
|
||||
_GEMINI_IMAGE_SIZE_1K_THRESHOLD: Final = (512 + 1024) // 2
|
||||
_GEMINI_IMAGE_SIZE_2K_THRESHOLD: Final = (1024 + 2048) // 2
|
||||
_GEMINI_IMAGE_SIZE_4K_THRESHOLD: Final = (2048 + 4096) // 2
|
||||
|
||||
GEMINI_IMAGE_ASPECT_RATIOS: Final[dict[str, float]] = {
|
||||
"1:1": 1 / 1,
|
||||
"1:4": 1 / 4,
|
||||
|
|
@ -105,8 +124,8 @@ def map_openai_size_to_gemini_image_config(size: str, model: str) -> dict[str, s
|
|||
return None
|
||||
|
||||
width, height = dimensions
|
||||
image_config: Final = {"aspectRatio": _map_dimensions_to_gemini_aspect_ratio(width, height)}
|
||||
image_size: Final = _map_dimensions_to_gemini_image_size(width, height)
|
||||
image_config: Final = {"aspectRatio": _map_dimensions_to_gemini_aspect_ratio(width, height, model)}
|
||||
image_size: Final = _map_dimensions_to_gemini_image_size(width, height, model)
|
||||
if is_gemini_image_model(model):
|
||||
if supports_gemini_image_size(model):
|
||||
image_config["imageSize"] = image_size
|
||||
|
|
@ -123,7 +142,7 @@ def supports_gemini_image_size(model: str) -> bool:
|
|||
return bool(value)
|
||||
except Exception:
|
||||
pass
|
||||
return "2.5-flash" not in model
|
||||
return _GEMINI_2_5_FLASH_MODEL_FAMILY not in model
|
||||
|
||||
|
||||
def is_gemini_image_model(model: str) -> bool:
|
||||
|
|
@ -333,24 +352,36 @@ def _parse_openai_image_size(size: str) -> tuple[int, int] | None:
|
|||
return width, height
|
||||
|
||||
|
||||
def _map_dimensions_to_gemini_aspect_ratio(width: int, height: int) -> str:
|
||||
if (width, height) in GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO:
|
||||
return GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO[(width, height)]
|
||||
def _map_dimensions_to_gemini_aspect_ratio(width: int, height: int, model: str) -> str:
|
||||
aspect_ratios: Final = tuple(
|
||||
ratio
|
||||
for ratio in GEMINI_IMAGE_ASPECT_RATIOS
|
||||
if _GEMINI_2_5_FLASH_MODEL_FAMILY not in model or ratio not in _GEMINI_2_5_FLASH_UNSUPPORTED_ASPECT_RATIOS
|
||||
)
|
||||
known_ratio: Final = GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO.get((width, height))
|
||||
if known_ratio is not None and known_ratio in aspect_ratios:
|
||||
return known_ratio
|
||||
|
||||
requested_ratio: Final = width / height
|
||||
return min(
|
||||
GEMINI_IMAGE_ASPECT_RATIOS,
|
||||
aspect_ratios,
|
||||
key=lambda aspect_ratio: abs(math.log(GEMINI_IMAGE_ASPECT_RATIOS[aspect_ratio] / requested_ratio)),
|
||||
)
|
||||
|
||||
|
||||
def _map_dimensions_to_gemini_image_size(width: int, height: int) -> str:
|
||||
effective_square_side: Final = math.sqrt(width * height)
|
||||
if effective_square_side < 768:
|
||||
return "512"
|
||||
if effective_square_side < 1536:
|
||||
def _map_dimensions_to_gemini_image_size(width: int, height: int, model: str) -> str:
|
||||
if _GEMINI_1K_ONLY_IMAGE_SIZE_MODEL in model:
|
||||
return "1K"
|
||||
if effective_square_side < 3072:
|
||||
|
||||
effective_square_side: Final = math.sqrt(width * height)
|
||||
if effective_square_side < _GEMINI_IMAGE_SIZE_1K_THRESHOLD:
|
||||
if is_gemini_image_model(model) and _GEMINI_512_IMAGE_SIZE_MODEL not in model:
|
||||
return "1K"
|
||||
else:
|
||||
return "512"
|
||||
if effective_square_side < _GEMINI_IMAGE_SIZE_2K_THRESHOLD:
|
||||
return "1K"
|
||||
if effective_square_side < _GEMINI_IMAGE_SIZE_4K_THRESHOLD:
|
||||
return "2K"
|
||||
return "4K"
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import math
|
||||
import re
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from copy import deepcopy
|
||||
|
|
@ -18,6 +19,11 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
|||
from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs
|
||||
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.gemini.common_utils import (
|
||||
GEMINI_IMAGE_ASPECT_RATIOS,
|
||||
GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO,
|
||||
map_openai_size_to_gemini_image_config,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.llms.vertex_ai import (
|
||||
VERTEX_AI_PROVIDER_METADATA_FIELDS,
|
||||
|
|
@ -35,6 +41,31 @@ GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType(
|
|||
GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/"
|
||||
|
||||
|
||||
# https://cloud.google.com/vertex-ai/generative-ai/docs/image/configure-aspect-ratio
|
||||
_IMAGEN_ASPECT_RATIOS: Final = MappingProxyType(
|
||||
{"1:1": 1.0, "3:4": 3 / 4, "4:3": 4 / 3, "9:16": 9 / 16, "16:9": 16 / 9}
|
||||
)
|
||||
|
||||
|
||||
def map_vertex_image_size(size: object, model: str, *, imagen: bool = False) -> Mapping[str, str] | None:
|
||||
if not isinstance(size, str) or (image_config := map_openai_size_to_gemini_image_config(size, model)) is None:
|
||||
return None
|
||||
|
||||
width, height = map(int, size.lower().split("x"))
|
||||
requested_ratio: Final = width / height
|
||||
aspect_ratio: Final = (
|
||||
min(_IMAGEN_ASPECT_RATIOS, key=lambda ratio: abs(math.log(_IMAGEN_ASPECT_RATIOS[ratio] / requested_ratio)))
|
||||
if imagen
|
||||
else image_config["aspectRatio"]
|
||||
)
|
||||
if (
|
||||
not math.isclose(GEMINI_IMAGE_ASPECT_RATIOS[aspect_ratio], requested_ratio)
|
||||
and GEMINI_IMAGE_SIZE_TO_ASPECT_RATIO.get((width, height)) != aspect_ratio
|
||||
):
|
||||
verbose_logger.warning("Vertex AI size %s is approximated with aspectRatio %s", size, aspect_ratio)
|
||||
return MappingProxyType({"aspectRatio": aspect_ratio} if imagen else image_config)
|
||||
|
||||
|
||||
def gemini_video_metadata_from_openai(video_metadata: Mapping[str, object]) -> dict[str, object]:
|
||||
return {
|
||||
gemini_key: video_metadata[openai_key]
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from httpx._types import RequestFiles
|
|||
import litellm
|
||||
from litellm.images.utils import ImageEditRequestUtils
|
||||
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.common_utils import get_vertex_base_url, map_vertex_image_size
|
||||
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
|
||||
|
|
@ -67,7 +67,7 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
return {}
|
||||
|
||||
size: Final = image_edit_optional_params.get("size")
|
||||
return {"aspectRatio": self._map_size_to_aspect_ratio(size or "")}
|
||||
return dict(map_vertex_image_size(size, model) or {"aspectRatio": "1:1"})
|
||||
|
||||
def _resolve_vertex_project(self) -> str | None:
|
||||
return (
|
||||
|
|
@ -176,11 +176,11 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
contents: Final[dict[str, object]] = {"role": "USER", "parts": parts}
|
||||
|
||||
# Add image-specific configuration
|
||||
image_config: Final = (
|
||||
{"aspect_ratio": image_edit_optional_request_params["aspectRatio"]}
|
||||
if "aspectRatio" in image_edit_optional_request_params
|
||||
else None
|
||||
)
|
||||
image_config: Final = {
|
||||
field: image_edit_optional_request_params[param]
|
||||
for param, field in (("aspectRatio", "aspect_ratio"), ("imageSize", "image_size"))
|
||||
if param in image_edit_optional_request_params
|
||||
}
|
||||
|
||||
generation_config: Final[dict[str, object]] = {
|
||||
key: value for key, value in (("response_modalities", ["IMAGE"]), ("image_config", image_config)) if value
|
||||
|
|
@ -222,17 +222,6 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
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 Gemini 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_inline_image_parts(self, image: FileTypes | list[FileTypes]) -> list[HttpxPartType]:
|
||||
images: Final[list[FileTypes]] = image if isinstance(image, list) else [image]
|
||||
return [
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ import litellm
|
|||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
from litellm.litellm_core_utils.url_utils import join_api_base
|
||||
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.common_utils import get_vertex_base_url, map_vertex_image_size
|
||||
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
|
||||
|
|
@ -70,13 +70,11 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
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
|
||||
size_params: Final = map_vertex_image_size(filtered_params.get("size"), model, imagen=True) or {}
|
||||
return {**mapped_params, **size_params}
|
||||
|
||||
def _resolve_vertex_project(self) -> str | None:
|
||||
return (
|
||||
|
|
@ -189,6 +187,11 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
"sampleCount": sample_count,
|
||||
"editMode": edit_mode,
|
||||
"editConfig": {"baseSteps": base_steps},
|
||||
**(
|
||||
{"aspectRatio": image_edit_optional_request_params["aspectRatio"]}
|
||||
if "aspectRatio" in image_edit_optional_request_params
|
||||
else {}
|
||||
),
|
||||
}
|
||||
|
||||
# Set default values for Vertex AI-specific parameters (not configurable by users via OpenAI API)
|
||||
|
|
@ -239,17 +242,6 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
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],
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from litellm.llms.gemini.common_utils import (
|
|||
get_gemini_image_web_search_requests,
|
||||
map_gemini_image_tools_params,
|
||||
)
|
||||
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
|
||||
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url, map_vertex_image_size
|
||||
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
|
||||
|
|
@ -89,8 +89,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
|
|||
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)
|
||||
continue
|
||||
elif k in ("aspectRatio", "aspect_ratio"):
|
||||
mapped_params["aspectRatio"] = v
|
||||
elif k in ("imageSize", "image_size"):
|
||||
|
|
@ -103,22 +102,23 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
|
|||
elif k not in ("tools", "web_search_options", "imageConfig"):
|
||||
mapped_params[k] = v
|
||||
|
||||
if "size" in non_default_params and "size" not in optional_params:
|
||||
image_params: Final = _JSON_OBJECT.validate_python(mapped_params)
|
||||
image_config: Final = _JSON_OBJECT.validate_python(image_params.get("imageConfig", {}))
|
||||
size_params: Final = {
|
||||
key: value
|
||||
for key, value in (
|
||||
map_vertex_image_size(non_default_params["size"], model) or {"aspectRatio": "1:1"}
|
||||
).items()
|
||||
if key not in image_config
|
||||
}
|
||||
return map_gemini_image_tools_params(
|
||||
_JSON_OBJECT.validate_python(non_default_params), {**size_params, **image_params}
|
||||
)
|
||||
|
||||
mapped_params = map_gemini_image_tools_params(non_default_params, mapped_params)
|
||||
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: 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 _resolve_vertex_project(self) -> str | None:
|
||||
return (
|
||||
getattr(self, "_vertex_project", None)
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from litellm.llms.base_llm.image_generation.transformation import (
|
|||
)
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
get_vertex_base_url,
|
||||
map_vertex_image_size,
|
||||
pop_vertex_request_labels,
|
||||
)
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
|
||||
|
|
@ -72,25 +73,14 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
|
|||
mapped_params["sampleCount"] = v
|
||||
elif k == "size":
|
||||
# Map OpenAI size format to Imagen aspectRatio
|
||||
mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(v)
|
||||
mapped_params["aspectRatio"] = (
|
||||
map_vertex_image_size(v, model, imagen=True) or {"aspectRatio": "1:1"}
|
||||
)["aspectRatio"]
|
||||
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: 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 _resolve_vertex_project(self) -> str | None:
|
||||
return (
|
||||
getattr(self, "_vertex_project", None)
|
||||
|
|
|
|||
|
|
@ -512,6 +512,40 @@ def test_map_openai_image_params_to_gemini_forwards_a_json_object_image_config()
|
|||
assert _map_image_config('{"aspectRatio": "1:1"}') == {"imageConfig": {"aspectRatio": "1:1"}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", ("", "gemini/", "vertex_ai/"))
|
||||
@pytest.mark.parametrize(
|
||||
"model, size, expected",
|
||||
(
|
||||
("gemini-3-pro-image", "256x256", {"aspectRatio": "1:1", "imageSize": "1K"}),
|
||||
("gemini-3-pro-image", "512x512", {"aspectRatio": "1:1", "imageSize": "1K"}),
|
||||
("gemini-3-pro-image-preview", "512x512", {"aspectRatio": "1:1", "imageSize": "1K"}),
|
||||
("gemini-3-pro-image", "2048x2048", {"aspectRatio": "1:1", "imageSize": "2K"}),
|
||||
("gemini-3-pro-image", "4096x4096", {"aspectRatio": "1:1", "imageSize": "4K"}),
|
||||
("gemini-3.1-flash-image", "512x512", {"aspectRatio": "1:1", "imageSize": "512"}),
|
||||
("gemini-3.1-flash-image-preview", "512x512", {"aspectRatio": "1:1", "imageSize": "512"}),
|
||||
("gemini-3.1-flash-image", "2048x2048", {"aspectRatio": "1:1", "imageSize": "2K"}),
|
||||
("gemini-3.1-flash-image", "4096x4096", {"aspectRatio": "1:1", "imageSize": "4K"}),
|
||||
("gemini-3.1-flash-lite-image", "256x256", {"aspectRatio": "1:1", "imageSize": "1K"}),
|
||||
("gemini-3.1-flash-lite-image", "512x512", {"aspectRatio": "1:1", "imageSize": "1K"}),
|
||||
("gemini-3.1-flash-lite-image", "2048x2048", {"aspectRatio": "1:1", "imageSize": "1K"}),
|
||||
("gemini-3.1-flash-lite-image", "4096x4096", {"aspectRatio": "1:1", "imageSize": "1K"}),
|
||||
("gemini-2.5-flash-image", "512x512", {"aspectRatio": "1:1"}),
|
||||
("gemini-2.5-flash-image", "1024x256", {"aspectRatio": "21:9"}),
|
||||
("gemini-2.5-flash-image", "256x1024", {"aspectRatio": "9:16"}),
|
||||
("gemini-2.5-flash-image", "8000x1000", {"aspectRatio": "21:9"}),
|
||||
("gemini-2.5-flash-image", "1000x8000", {"aspectRatio": "9:16"}),
|
||||
("gemini-3.1-flash-image", "1024x256", {"aspectRatio": "4:1", "imageSize": "512"}),
|
||||
),
|
||||
)
|
||||
def test_size_mapping_respects_model_image_limits(
|
||||
prefix: str, model: str, size: str, expected: dict[str, str]
|
||||
) -> None:
|
||||
"""Limits as of 2026-10-26, https://ai.google.dev/gemini-api/docs/image-generation#aspect_ratios_and_image_size."""
|
||||
assert map_openai_image_params_to_gemini({"size": size}, prefix + model, ["size"]) == {
|
||||
"imageConfig": expected
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw_image_config", ["[1, 2]", '"16:9"', "null"])
|
||||
def test_map_openai_image_params_to_gemini_drops_or_refuses_a_non_object_image_config(raw_image_config: str):
|
||||
with contextlib.suppress(litellm.UnsupportedParamsError):
|
||||
|
|
|
|||
|
|
@ -8289,11 +8289,12 @@ def test_gemini_image_generation_ignores_non_dict_image_config():
|
|||
def test_gemini_image_generation_openai_size_maps_to_google_table(
|
||||
size: str, expected_aspect_ratio: str, expected_image_size: str
|
||||
):
|
||||
"""Flash table as of 2026-10-26, https://ai.google.dev/gemini-api/docs/image-generation#aspect_ratios_and_image_size."""
|
||||
from litellm.llms.gemini.common_utils import (
|
||||
map_openai_size_to_gemini_image_config,
|
||||
)
|
||||
|
||||
assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") == {
|
||||
assert map_openai_size_to_gemini_image_config(size, "gemini-3.1-flash-image") == {
|
||||
"aspectRatio": expected_aspect_ratio,
|
||||
"imageSize": expected_image_size,
|
||||
}
|
||||
|
|
@ -8317,7 +8318,7 @@ def test_gemini_image_generation_openai_size_snaps_to_nearest_option(
|
|||
map_openai_size_to_gemini_image_config,
|
||||
)
|
||||
|
||||
assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") == {
|
||||
assert map_openai_size_to_gemini_image_config(size, "gemini-3.1-flash-image") == {
|
||||
"aspectRatio": expected_aspect_ratio,
|
||||
"imageSize": expected_image_size,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,22 +3,118 @@ import json
|
|||
import os
|
||||
from datetime import datetime
|
||||
from io import BytesIO
|
||||
from typing import Dict
|
||||
from typing import Dict, Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.gemini.common_utils import map_openai_size_to_gemini_image_config
|
||||
from litellm.llms.vertex_ai.image_edit.vertex_gemini_transformation import (
|
||||
VertexAIGeminiImageEditConfig,
|
||||
)
|
||||
from litellm.llms.vertex_ai.image_edit.vertex_imagen_transformation import (
|
||||
VertexAIImagenImageEditConfig,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size", ("1536x1024", "1024x1536", "3072x2048", "256x256", "512x512", "4096x4096"))
|
||||
@pytest.mark.parametrize("model", ("gemini-3.1-flash-image", "gemini-3-pro-image", "gemini-3.1-flash-lite-image"))
|
||||
def test_vertex_gemini_edit_size_reaches_request(size: str, model: str) -> None:
|
||||
config: Final = VertexAIGeminiImageEditConfig()
|
||||
payload, _ = config.transform_image_edit_request(
|
||||
model,
|
||||
"mountains",
|
||||
BytesIO(b"image"),
|
||||
config.map_openai_params({"size": size}, model, False),
|
||||
GenericLiteLLMParams(),
|
||||
{},
|
||||
)
|
||||
request: Final = TypeAdapter(dict[str, JsonValue]).validate_json(TypeAdapter(str).validate_python(payload))
|
||||
expected: Final = map_openai_size_to_gemini_image_config(size, model) or {}
|
||||
assert request["generationConfig"] == {
|
||||
"response_modalities": ["IMAGE"],
|
||||
**(
|
||||
{"image_config": {"aspect_ratio": expected["aspectRatio"], "image_size": expected["imageSize"]}}
|
||||
if expected
|
||||
else {}
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"size, expected_ratio",
|
||||
(("1536x1024", "4:3"), ("1024x1536", "3:4"), ("100000x1", "16:9")),
|
||||
)
|
||||
def test_vertex_imagen_edit_size_reaches_request(size: str, expected_ratio: str | None) -> None:
|
||||
model: Final = "imagen-3.0-capability-001"
|
||||
config: Final = VertexAIImagenImageEditConfig()
|
||||
payload, _ = config.transform_image_edit_request(
|
||||
model,
|
||||
"mountains",
|
||||
BytesIO(b"image"),
|
||||
config.map_openai_params({"size": size, "n": 2}, model, False),
|
||||
GenericLiteLLMParams(),
|
||||
{},
|
||||
)
|
||||
request: Final = TypeAdapter(dict[str, JsonValue]).validate_json(TypeAdapter(str).validate_python(payload))
|
||||
assert request["parameters"] == {
|
||||
"sampleCount": 2,
|
||||
"editMode": "EDIT_MODE_INPAINT_INSERTION",
|
||||
"editConfig": {"baseSteps": 50},
|
||||
"guidanceScale": 7.5,
|
||||
"seed": None,
|
||||
**({"aspectRatio": expected_ratio} if expected_ratio is not None else {}),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size", ("auto", "unknown", "0x1024", "-1x1024", "1024x", "3:2"))
|
||||
@pytest.mark.parametrize(
|
||||
"config_type, model, config_key, expected",
|
||||
(
|
||||
(
|
||||
VertexAIGeminiImageEditConfig,
|
||||
"gemini-3.1-flash-image",
|
||||
"generationConfig",
|
||||
{"response_modalities": ["IMAGE"], "image_config": {"aspect_ratio": "1:1"}},
|
||||
),
|
||||
(
|
||||
VertexAIImagenImageEditConfig,
|
||||
"imagen-3.0-capability-001",
|
||||
"parameters",
|
||||
{
|
||||
"sampleCount": 1,
|
||||
"editMode": "EDIT_MODE_INPAINT_INSERTION",
|
||||
"editConfig": {"baseSteps": 50},
|
||||
"guidanceScale": 7.5,
|
||||
"seed": None,
|
||||
},
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_vertex_edit_preserves_legacy_size_fallback(
|
||||
size: str,
|
||||
config_type: type[VertexAIGeminiImageEditConfig] | type[VertexAIImagenImageEditConfig],
|
||||
model: str,
|
||||
config_key: str,
|
||||
expected: dict[str, object],
|
||||
) -> None:
|
||||
config: Final = config_type()
|
||||
payload, _ = config.transform_image_edit_request(
|
||||
model,
|
||||
"mountains",
|
||||
BytesIO(b"image"),
|
||||
config.map_openai_params({"size": size}, model, False),
|
||||
GenericLiteLLMParams(),
|
||||
{},
|
||||
)
|
||||
request: Final = TypeAdapter(dict[str, JsonValue]).validate_json(TypeAdapter(str).validate_python(payload))
|
||||
assert request[config_key] == expected
|
||||
|
||||
|
||||
class TestVertexAIGeminiImageEditTransformation:
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from pydantic import ValidationError
|
|||
import litellm
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.gemini.image_generation.transformation import GoogleImageGenConfig
|
||||
from litellm.llms.vertex_ai.image_generation import (
|
||||
get_vertex_ai_image_generation_config,
|
||||
)
|
||||
|
|
@ -23,6 +24,98 @@ from litellm.llms.vertex_ai.image_generation.vertex_imagen_transformation import
|
|||
from litellm.types.utils import ImageResponse
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size", ("1536x1024", "1024x1536", "3072x2048", "256x256", "512x512", "4096x4096"))
|
||||
@pytest.mark.parametrize("model", ("gemini-3.1-flash-image", "gemini-3-pro-image", "gemini-3.1-flash-lite-image"))
|
||||
def test_vertex_gemini_size_matches_direct_gemini_request(size: str, model: str) -> None:
|
||||
vertex: Final = VertexAIGeminiImageGenerationConfig()
|
||||
direct: Final = GoogleImageGenConfig()
|
||||
vertex_request: Final = vertex.transform_image_generation_request(
|
||||
model, "mountains", vertex.map_openai_params({"size": size}, {}, model, False), {}, {}
|
||||
)
|
||||
direct_request: Final = direct.transform_image_generation_request(
|
||||
model, "mountains", direct.map_openai_params({"size": size}, {}, model, False), {}, {}
|
||||
)
|
||||
assert vertex_request["generationConfig"].get("imageConfig") == direct_request["generationConfig"].get(
|
||||
"imageConfig"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"size, expected_ratio",
|
||||
(("1536x1024", "4:3"), ("1024x1536", "3:4"), ("100000x1", "16:9")),
|
||||
)
|
||||
def test_vertex_imagen_generation_size_reaches_request(size: str, expected_ratio: str | None) -> None:
|
||||
model: Final = "imagen-4.0-generate-001"
|
||||
config: Final = VertexAIImagenImageGenerationConfig()
|
||||
request: Final = config.transform_image_generation_request(
|
||||
model, "mountains", config.map_openai_params({"size": size, "n": 2}, {}, model, False), {}, {}
|
||||
)
|
||||
assert request["parameters"] == {
|
||||
"sampleCount": 2,
|
||||
**({"aspectRatio": expected_ratio} if expected_ratio is not None else {}),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size", ("auto", "unknown", "0x1024", "-1x1024", "1024x", "3:2"))
|
||||
@pytest.mark.parametrize(
|
||||
"config_type, model, config_key, expected",
|
||||
(
|
||||
(
|
||||
VertexAIGeminiImageGenerationConfig,
|
||||
"gemini-3.1-flash-image",
|
||||
"generationConfig",
|
||||
{"responseModalities": ["IMAGE"], "imageConfig": {"aspectRatio": "1:1"}},
|
||||
),
|
||||
(
|
||||
VertexAIImagenImageGenerationConfig,
|
||||
"imagen-4.0-generate-001",
|
||||
"parameters",
|
||||
{"sampleCount": 1, "aspectRatio": "1:1"},
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_vertex_generation_preserves_legacy_size_fallback(
|
||||
size: str,
|
||||
config_type: type[VertexAIGeminiImageGenerationConfig] | type[VertexAIImagenImageGenerationConfig],
|
||||
model: str,
|
||||
config_key: str,
|
||||
expected: dict[str, object],
|
||||
) -> None:
|
||||
config: Final = config_type()
|
||||
request: Final = config.transform_image_generation_request(
|
||||
model, "mountains", config.map_openai_params({"size": size}, {}, model, False), {}, {}
|
||||
)
|
||||
assert request[config_key] == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("config_type", (VertexAIGeminiImageGenerationConfig, VertexAIImagenImageGenerationConfig))
|
||||
def test_vertex_generation_warns_when_size_ratio_is_approximated(
|
||||
config_type: type[VertexAIGeminiImageGenerationConfig] | type[VertexAIImagenImageGenerationConfig],
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
mapped: Final = config_type().map_openai_params({"size": "100000x1"}, {}, "gemini-3.1-flash-image", False)
|
||||
assert "100000x1" in caplog.text
|
||||
assert str(mapped["aspectRatio"]) in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size_first", (True, False))
|
||||
@pytest.mark.parametrize(
|
||||
"overrides",
|
||||
(
|
||||
{"aspect_ratio": "9:16", "image_size": "2K"},
|
||||
{"imageConfig": {"aspectRatio": "9:16", "imageSize": "2K"}},
|
||||
),
|
||||
)
|
||||
def test_vertex_gemini_explicit_image_settings_override_size(size_first: bool, overrides: dict[str, object]) -> None:
|
||||
model: Final = "gemini-3.1-flash-image"
|
||||
config: Final = VertexAIGeminiImageGenerationConfig()
|
||||
params: Final = {"size": "1536x1024", **overrides} if size_first else {**overrides, "size": "1536x1024"}
|
||||
request: Final = config.transform_image_generation_request(
|
||||
model, "mountains", config.map_openai_params(params, {}, model, False), {}, {}
|
||||
)
|
||||
assert request["generationConfig"]["imageConfig"] == {"aspectRatio": "9:16", "imageSize": "2K"}
|
||||
|
||||
|
||||
class TestVertexAIGeminiImageGenerationConfig:
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures"""
|
||||
|
|
@ -55,14 +148,13 @@ class TestVertexAIGeminiImageGenerationConfig:
|
|||
result = self.config.map_openai_params(non_default_params, optional_params, "gemini-2.5-flash-image", False)
|
||||
assert result.get("aspectRatio") == "16:9"
|
||||
|
||||
def test_map_size_to_aspect_ratio(self):
|
||||
"""Test size to aspect ratio mapping"""
|
||||
assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1"
|
||||
assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9"
|
||||
assert self.config._map_size_to_aspect_ratio("1024x1792") == "9:16"
|
||||
assert self.config._map_size_to_aspect_ratio("1280x896") == "4:3"
|
||||
assert self.config._map_size_to_aspect_ratio("896x1280") == "3:4"
|
||||
assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default
|
||||
@pytest.mark.parametrize(
|
||||
"size, expected_ratio",
|
||||
(("1024x1024", "1:1"), ("1792x1024", "16:9"), ("1024x1792", "9:16"), ("1280x896", "4:3"), ("896x1280", "3:4")),
|
||||
)
|
||||
def test_map_openai_params_legacy_sizes(self, size: str, expected_ratio: str) -> None:
|
||||
mapped: Final = self.config.map_openai_params({"size": size}, {}, "gemini-3.1-flash-image", False)
|
||||
assert mapped["aspectRatio"] == expected_ratio
|
||||
|
||||
def test_get_supported_openai_params_includes_native_gemini_params(self):
|
||||
"""Test that native Gemini imageConfig params are supported"""
|
||||
|
|
@ -474,11 +566,10 @@ class TestVertexAIImagenImageGenerationConfig:
|
|||
result = self.config.map_openai_params(non_default_params, optional_params, "imagegeneration@006", False)
|
||||
assert result.get("aspectRatio") == "1:1"
|
||||
|
||||
def test_map_size_to_aspect_ratio(self):
|
||||
"""Test size to aspect ratio mapping"""
|
||||
assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1"
|
||||
assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9"
|
||||
assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default
|
||||
@pytest.mark.parametrize("size, expected_ratio", (("1024x1024", "1:1"), ("1792x1024", "16:9")))
|
||||
def test_map_openai_params_legacy_sizes(self, size: str, expected_ratio: str) -> None:
|
||||
mapped: Final = self.config.map_openai_params({"size": size}, {}, "imagen-4.0-generate-001", False)
|
||||
assert mapped["aspectRatio"] == expected_ratio
|
||||
|
||||
def test_transform_image_generation_request_basic(self):
|
||||
"""Test basic request transformation"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue