Add Gemini image edit support (#16430)

* Add gemini image edit support * fix lint errors * fix lint errors * fix lint errors * Add docs
2026-08-23 18:25:22 +00:00 · 2025-11-12 18:48:27 -08:00
parent 8bf491c939
commit 018bd2e039
10 changed files with 493 additions and 4 deletions
@@ -14,9 +14,9 @@ LiteLLM provides image editing functionality that maps to OpenAI's `/images/edit
 | Fallbacks | ✅ | Works between supported models |
 | Loadbalancing | ✅ | Works between supported models |
 | Supported operations | Create image edits | Single and multiple images supported |
-| Supported LiteLLM SDK Versions | 1.63.8+ | |
-| Supported LiteLLM Proxy Versions | 1.71.1+ | |
-| Supported LLM providers | **OpenAI** | Currently only `openai` is supported |
+| Supported LiteLLM SDK Versions | 1.63.8+ | Gemini support requires 1.79.3+ |
+| Supported LiteLLM Proxy Versions | 1.71.1+ | Gemini support requires 1.79.3+ |
+| Supported LLM providers | **OpenAI**, **Gemini (Google AI Studio)** | Gemini supports the new `gemini-2.5-flash-image` family |

 #### ⚡️See all supported models and providers at [models.litellm.ai](https://models.litellm.ai/)

@@ -149,6 +149,54 @@ for i, image_data in enumerate(response.data):
    print(f"Image {i+1}: {image_data.url}")
 ```

+```
+
+</TabItem>
+
+<TabItem value="gemini" label="Gemini">
+
+#### Basic Image Edit
+```python showLineNumbers title="Gemini Image Edit"
+import base64
+import os
+from litellm import image_edit
+
+os.environ["GEMINI_API_KEY"] = "your-api-key"
+
+response = image_edit(
+    model="gemini/gemini-2.5-flash-image",
+    image=open("original_image.png", "rb"),
+    prompt="Add aurora borealis to the night sky",
+    size="1792x1024",  # mapped to aspectRatio=16:9 for Gemini
+)
+
+edited_image_bytes = base64.b64decode(response.data[0].b64_json)
+with open("edited_image.png", "wb") as f:
+    f.write(edited_image_bytes)
+```
+
+#### Multiple Images Edit
+```python showLineNumbers title="Gemini Multiple Images Edit"
+import base64
+import os
+from litellm import image_edit
+
+os.environ["GEMINI_API_KEY"] = "your-api-key"
+
+response = image_edit(
+    model="gemini/gemini-2.5-flash-image",
+    image=[
+        open("scene.png", "rb"),
+        open("style_reference.png", "rb"),
+    ],
+    prompt="Blend the reference style into the scene while keeping the subject sharp.",
+)
+
+for idx, image_obj in enumerate(response.data):
+    with open(f"gemini_edit_{idx}.png", "wb") as f:
+        f.write(base64.b64decode(image_obj.b64_json))
+```
+
 </TabItem>
 </Tabs>

@@ -224,6 +272,36 @@ curl -X POST "http://localhost:4000/v1/images/edits" \
  -F "response_format=url"
 ```

+```
+
+</TabItem>
+
+<TabItem value="gemini" label="Gemini">
+
+1. Add the Gemini image edit model to your `config.yaml`:
+```yaml showLineNumbers title="Gemini Proxy Configuration"
+model_list:
+  - model_name: gemini-image-edit
+    litellm_params:
+      model: gemini/gemini-2.5-flash-image
+      api_key: os.environ/GEMINI_API_KEY
+```
+
+2. Start the LiteLLM proxy server:
+```bash showLineNumbers title="Start LiteLLM Proxy Server"
+litellm --config /path/to/config.yaml
+```
+
+3. Make an image edit request (Gemini responses are base64-only):
+```bash showLineNumbers title="Gemini Proxy Image Edit"
+curl -X POST "http://0.0.0.0:4000/v1/images/edits" \
+  -H "Authorization: Bearer <YOUR-LITELLM-KEY>" \
+  -F "model=gemini-image-edit" \
+  -F "image=@original_image.png" \
+  -F "prompt=Add a warm golden-hour glow to the scene" \
+  -F "size=1024x1024"
+```
+
 </TabItem>
 </Tabs>

@@ -10,7 +10,7 @@ import TabItem from '@theme/TabItem';
 | Provider Route on LiteLLM | `gemini/` |
 | Provider Doc | [Google AI Studio ↗](https://aistudio.google.com/) |
 | API Endpoint for Provider | https://generativelanguage.googleapis.com |
-| Supported OpenAI Endpoints | `/chat/completions`, [`/embeddings`](../embedding/supported_embedding#gemini-ai-embedding-models), `/completions`, [`/videos`](./gemini/videos.md) |
+| Supported OpenAI Endpoints | `/chat/completions`, [`/embeddings`](../embedding/supported_embedding#gemini-ai-embedding-models), `/completions`, [`/videos`](./gemini/videos.md), [`/images/edits`](../image_edits.md) |
 | Pass-through Endpoint | [Supported](../pass_through/google_ai_studio.md) |

 <br />
@@ -943,6 +943,7 @@ def completion_cost(  # noqa: PLR0915
                        n=n,
                        size=size,
                        optional_params=optional_params,
+                        call_type=call_type,
                    )
                elif (
                    call_type == CallTypes.create_video.value
@@ -640,6 +640,7 @@ class CostCalculatorUtils:
        n: Optional[int] = None,
        size: Optional[str] = None,
        optional_params: Optional[dict] = None,
+        call_type: Optional[str] = None,
    ) -> float:
        """
        Route the image generation cost calculator based on the custom_llm_provider
@@ -713,6 +714,18 @@ class CostCalculatorUtils:
                image_response=completion_response,
            )
        elif custom_llm_provider == litellm.LlmProviders.GEMINI.value:
+            if call_type in (
+                CallTypes.image_edit.value,
+                CallTypes.aimage_edit.value,
+            ):
+                from litellm.llms.gemini.image_edit.cost_calculator import (
+                    cost_calculator as gemini_image_edit_cost_calculator,
+                )
+
+                return gemini_image_edit_cost_calculator(
+                    model=model,
+                    image_response=completion_response,
+                )
            from litellm.llms.gemini.image_generation.cost_calculator import (
                cost_calculator as gemini_image_cost_calculator,
            )
@@ -0,0 +1,11 @@
+from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
+
+from .transformation import GeminiImageEditConfig
+from .cost_calculator import cost_calculator
+
+__all__ = ["GeminiImageEditConfig", "get_gemini_image_edit_config", "cost_calculator"]
+
+
+def get_gemini_image_edit_config(model: str) -> BaseImageEditConfig:
+    return GeminiImageEditConfig()
+
@@ -0,0 +1,35 @@
+"""
+Gemini Image Edit Cost Calculator
+"""
+
+from typing import Any
+
+import litellm
+from litellm.types.utils import ImageResponse
+
+
+def cost_calculator(
+    model: str,
+    image_response: Any,
+) -> float:
+    """
+    Gemini image edit cost calculator.
+
+    Mirrors image generation pricing: charge per returned image based on
+    model metadata (`output_cost_per_image`).
+    """
+    model_info = litellm.get_model_info(
+        model=model,
+        custom_llm_provider="gemini",
+    )
+
+    output_cost_per_image: float = model_info.get("output_cost_per_image") or 0.0
+
+    if not isinstance(image_response, ImageResponse):
+        raise ValueError(
+            f"image_response must be of type ImageResponse got type={type(image_response)}"
+        )
+
+    num_images = len(image_response.data or [])
+    return output_cost_per_image * num_images
+
@@ -0,0 +1,197 @@
+import base64
+from io import BufferedReader, BytesIO
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
+
+import httpx
+from httpx._types import RequestFiles
+
+from litellm.images.utils import ImageEditRequestUtils
+from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
+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 GeminiImageEditConfig(BaseImageEditConfig):
+    DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta"
+    SUPPORTED_PARAMS: List[str] = ["size"]
+
+    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 = 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: Dict[str, Any] = {}
+
+        if "size" in filtered_params:
+            mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(
+                filtered_params["size"]  # type: ignore[arg-type]
+            )
+
+        return mapped_params
+
+    def validate_environment(
+        self,
+        headers: dict,
+        model: str,
+        api_key: Optional[str] = None,
+    ) -> dict:
+        final_api_key: Optional[str] = api_key or get_secret_str("GEMINI_API_KEY")
+        if not final_api_key:
+            raise ValueError("GEMINI_API_KEY is not set")
+
+        headers["x-goog-api-key"] = final_api_key
+        headers["Content-Type"] = "application/json"
+        return headers
+
+    def get_complete_url(
+        self,
+        model: str,
+        api_base: Optional[str],
+        litellm_params: dict,
+    ) -> str:
+        base_url = api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL
+        base_url = base_url.rstrip("/")
+        return f"{base_url}/models/{model}:generateContent"
+
+    def transform_image_edit_request(  # type: ignore[override]
+        self,
+        model: str,
+        prompt: str,
+        image: FileTypes,
+        image_edit_optional_request_params: Dict[str, Any],
+        litellm_params: GenericLiteLLMParams,
+        headers: dict,
+    ) -> Tuple[Dict[str, Any], Optional[RequestFiles]]:
+        inline_parts = self._prepare_inline_image_parts(image)
+        if not inline_parts:
+            raise ValueError("Gemini image edit requires at least one image.")
+
+        contents = [
+            {
+                "parts": inline_parts + [{"text": prompt}],
+            }
+        ]
+
+        request_body: Dict[str, Any] = {"contents": contents}
+
+        generation_config: Dict[str, Any] = {}
+
+        if "aspectRatio" in image_edit_optional_request_params:
+            generation_config["aspectRatio"] = image_edit_optional_request_params[
+                "aspectRatio"
+            ]
+
+        if generation_config:
+            request_body["generationConfig"] = generation_config
+
+        empty_files = cast(RequestFiles, [])
+        return request_body, empty_files
+
+    def transform_image_edit_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        logging_obj: Any,
+    ) -> ImageResponse:
+        model_response = ImageResponse()
+        try:
+            response_json = 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,
+            )
+
+        candidates = response_json.get("candidates", [])
+        data_list: List[ImageObject] = []
+
+        for candidate in candidates:
+            content = candidate.get("content", {})
+            parts = content.get("parts", [])
+            for part in parts:
+                inline_data = part.get("inlineData")
+                if inline_data and inline_data.get("data"):
+                    data_list.append(
+                        ImageObject(
+                            b64_json=inline_data["data"],
+                            url=None,
+                        )
+                    )
+
+        model_response.data = cast(List[OpenAIImage], data_list)
+        return model_response
+
+    def _map_size_to_aspect_ratio(self, size: str) -> str:
+        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 _prepare_inline_image_parts(
+        self, image: Union[FileTypes, List[FileTypes]]
+    ) -> List[Dict[str, Any]]:
+        images: List[FileTypes]
+        if isinstance(image, list):
+            images = image
+        else:
+            images = [image]
+
+        inline_parts: List[Dict[str, Any]] = []
+        for img in images:
+            if img is None:
+                continue
+
+            mime_type = ImageEditRequestUtils.get_image_content_type(img)
+            image_bytes = self._read_all_bytes(img)
+            inline_parts.append(
+                {
+                    "inlineData": {
+                        "mimeType": mime_type,
+                        "data": base64.b64encode(image_bytes).decode("utf-8"),
+                    }
+                }
+            )
+
+        return inline_parts
+
+    def _read_all_bytes(self, image: FileTypes) -> bytes:
+        if isinstance(image, bytes):
+            return image
+        if isinstance(image, BytesIO):
+            current_pos = image.tell()
+            image.seek(0)
+            data = image.read()
+            image.seek(current_pos)
+            return data
+        if isinstance(image, BufferedReader):
+            current_pos = image.tell()
+            image.seek(0)
+            data = image.read()
+            image.seek(current_pos)
+            return data
+        raise ValueError("Unsupported image type for Gemini image edit.")
@@ -7720,6 +7720,10 @@ class ProviderConfigManager:
            from litellm.llms.azure_ai.image_edit import get_azure_ai_image_edit_config

            return get_azure_ai_image_edit_config(model)
+        elif LlmProviders.GEMINI == provider:
+            from litellm.llms.gemini.image_edit import get_gemini_image_edit_config
+
+            return get_gemini_image_edit_config(model)
        elif LlmProviders.LITELLM_PROXY == provider:
            from litellm.llms.litellm_proxy.image_edit.transformation import (
                LiteLLMProxyImageEditConfig,
@@ -11670,6 +11670,7 @@
        "litellm_provider": "vertex_ai-language-models",
        "max_audio_length_hours": 8.4,
        "max_audio_per_prompt": 1,
+        "supports_reasoning": false,
        "max_images_per_prompt": 3000,
        "max_input_tokens": 32768,
        "max_output_tokens": 32768,
@@ -0,0 +1,149 @@
+import base64
+import json
+from io import BytesIO
+from typing import Dict
+from unittest.mock import MagicMock
+
+import httpx
+import pytest
+
+from litellm.llms.gemini.image_edit.transformation import GeminiImageEditConfig
+
+
+class TestGeminiImageEditTransformation:
+    def setup_method(self) -> None:
+        self.config = GeminiImageEditConfig()
+        self.model = "gemini-2.5-flash-image-preview"
+        self.prompt = "Enhance this photo with a dramatic night sky."
+        self.logging_obj = MagicMock()
+
+    def test_map_openai_params(self) -> None:
+        optional_params: Dict[str, object] = {
+            "size": "1792x1024",
+            "response_format": "b64_json",
+            "quality": "high",
+        }
+
+        mapped = self.config.map_openai_params(
+            image_edit_optional_params=optional_params,  # type: ignore[arg-type]
+            model=self.model,
+            drop_params=False,
+        )
+
+        assert mapped["aspectRatio"] == "16:9"
+        assert "response_format" not in mapped
+        assert "quality" not in mapped
+
+    def test_transform_image_edit_request(self) -> None:
+        image_bytes = b"fake_image_data"
+        image = BytesIO(image_bytes)
+        optional_params = {
+            "sampleCount": 2,
+            "aspectRatio": "16:9",
+        }
+
+        request_body, files = self.config.transform_image_edit_request(
+            model=self.model,
+            prompt=self.prompt,
+            image=[image],  # Gemini pipeline passes list of images
+            image_edit_optional_request_params=optional_params,
+            litellm_params=MagicMock(),
+            headers={},
+        )
+
+        assert files == []
+
+        parts = request_body["contents"][0]["parts"]
+        assert parts[-1]["text"] == self.prompt
+
+        inline_data = parts[0]["inlineData"]
+        assert inline_data["mimeType"] == "image/png"
+        assert base64.b64decode(inline_data["data"]) == image_bytes
+
+        generation_config = request_body["generationConfig"]
+        assert generation_config["aspectRatio"] == "16:9"
+
+    def test_transform_image_edit_request_multiple_images(self) -> None:
+        image_one = BytesIO(b"image_one")
+        image_two = BytesIO(b"image_two")
+
+        request_body, files = self.config.transform_image_edit_request(
+            model=self.model,
+            prompt=self.prompt,
+            image=[image_one, image_two],
+            image_edit_optional_request_params={},
+            litellm_params=MagicMock(),
+            headers={},
+        )
+
+        assert files == []
+        parts = request_body["contents"][0]["parts"]
+
+        assert len(parts) == 3  # two images + text prompt
+        assert parts[-1]["text"] == self.prompt
+        assert base64.b64decode(parts[0]["inlineData"]["data"]) == b"image_one"
+        assert base64.b64decode(parts[1]["inlineData"]["data"]) == b"image_two"
+
+    def test_transform_image_edit_response(self) -> None:
+        response_payload = {
+            "candidates": [
+                {
+                    "content": {
+                        "parts": [
+                            {
+                                "inlineData": {
+                                    "mimeType": "image/png",
+                                    "data": base64.b64encode(b"image-one").decode("utf-8"),
+                                }
+                            }
+                        ]
+                    }
+                },
+                {
+                    "content": {
+                        "parts": [
+                            {
+                                "inlineData": {
+                                    "mimeType": "image/png",
+                                    "data": base64.b64encode(b"image-two").decode("utf-8"),
+                                }
+                            }
+                        ]
+                    }
+                },
+            ]
+        }
+
+        mock_response = MagicMock(spec=httpx.Response)
+        mock_response.json.return_value = response_payload
+        mock_response.status_code = 200
+        mock_response.headers = {}
+
+        image_response = self.config.transform_image_edit_response(
+            model=self.model,
+            raw_response=mock_response,
+            logging_obj=self.logging_obj,
+        )
+
+        assert image_response.data is not None
+        assert len(image_response.data) == 2
+        assert image_response.data[0].b64_json == base64.b64encode(b"image-one").decode(
+            "utf-8"
+        )
+        assert image_response.data[1].b64_json == base64.b64encode(b"image-two").decode(
+            "utf-8"
+        )
+
+    def test_transform_image_edit_request_without_image_raises(self) -> None:
+        optional_params = {}
+
+        with pytest.raises(ValueError):
+            self.config.transform_image_edit_request(
+                model=self.model,
+                prompt=self.prompt,
+                image=[],
+                image_edit_optional_request_params=optional_params,
+                litellm_params=MagicMock(),
+                headers={},
+            )
+