invoke-ai--invokeai/invokeai/app/services/external_generation/providers/gemini.py

from __future__ import annotations

import requests

from invokeai.app.services.external_generation.errors import (
    ExternalProviderRateLimitError,
    ExternalProviderRequestError,
)
from invokeai.app.services.external_generation.external_generation_base import ExternalProvider
from invokeai.app.services.external_generation.external_generation_common import (
    ExternalGeneratedImage,
    ExternalGenerationRequest,
    ExternalGenerationResult,
)
from invokeai.app.services.external_generation.image_utils import decode_image_base64, encode_image_base64


class GeminiProvider(ExternalProvider):
    provider_id = "gemini"
    _SYSTEM_INSTRUCTION = (
        "You are an image generation model. Always respond with an image based on the user's prompt. "
        "Do not return text-only responses. If the user input is not an edit instruction, "
        "interpret it as a request to create a new image."
    )

    def is_configured(self) -> bool:
        return bool(self._app_config.external_gemini_api_key)

    def generate(self, request: ExternalGenerationRequest) -> ExternalGenerationResult:
        api_key = self._app_config.external_gemini_api_key
        if not api_key:
            raise ExternalProviderRequestError("Gemini API key is not configured")

        base_url = (self._app_config.external_gemini_base_url or "https://generativelanguage.googleapis.com").rstrip(
            "/"
        )
        if not base_url.endswith("/v1") and not base_url.endswith("/v1beta"):
            base_url = f"{base_url}/v1beta"
        model_id = request.model.provider_model_id.removeprefix("models/")
        endpoint = f"{base_url}/models/{model_id}:generateContent"

        request_parts: list[dict[str, object]] = []

        if request.init_image is not None:
            request_parts.append(
                {
                    "inlineData": {
                        "mimeType": "image/png",
                        "data": encode_image_base64(request.init_image),
                    }
                }
            )

        request_parts.append({"text": request.prompt})

        for reference in request.reference_images:
            request_parts.append(
                {
                    "inlineData": {
                        "mimeType": "image/png",
                        "data": encode_image_base64(reference.image),
                    }
                }
            )

        opts = request.provider_options or {}

        generation_config: dict[str, object] = {
            "candidateCount": request.num_images,
            "responseModalities": ["IMAGE"],
        }
        if "temperature" in opts:
            generation_config["temperature"] = opts["temperature"]
        aspect_ratio = _select_aspect_ratio(
            request.width,
            request.height,
            request.model.capabilities.allowed_aspect_ratios,
        )
        uses_image_config = request.model.capabilities.resolution_presets is not None
        if uses_image_config:
            image_config: dict[str, str] = {}
            if aspect_ratio is not None:
                image_config["aspectRatio"] = aspect_ratio
            if request.image_size is not None:
                image_config["imageSize"] = request.image_size
            if image_config:
                generation_config["imageConfig"] = image_config
        system_instruction = self._SYSTEM_INSTRUCTION
        if request.init_image is not None:
            system_instruction = (
                f"{system_instruction} An input image is provided. "
                "Treat the prompt as an edit instruction and modify the image accordingly. "
                "Do not return the original image unchanged."
            )
        if not uses_image_config and aspect_ratio is not None:
            system_instruction = f"{system_instruction} Use an aspect ratio of {aspect_ratio}."

        payload: dict[str, object] = {
            "systemInstruction": {"parts": [{"text": system_instruction}]},
            "contents": [{"role": "user", "parts": request_parts}],
            "generationConfig": generation_config,
        }
        if "thinking_level" in opts:
            payload["thinkingConfig"] = {"thinkingLevel": opts["thinking_level"].upper()}

        response = requests.post(
            endpoint,
            params={"key": api_key},
            json=payload,
            timeout=120,
        )

        if not response.ok:
            if response.status_code == 429:
                retry_after = _parse_retry_after(response.headers.get("retry-after"))
                raise ExternalProviderRateLimitError(
                    f"Gemini rate limit exceeded. {f'Retry after {retry_after:.0f}s.' if retry_after else 'Please try again later.'}",
                    retry_after=retry_after,
                )
            raise ExternalProviderRequestError(
                f"Gemini request failed with status {response.status_code} for model '{model_id}': {response.text}"
            )

        data = response.json()
        if not isinstance(data, dict):
            raise ExternalProviderRequestError("Gemini response payload was not a JSON object")
        images: list[ExternalGeneratedImage] = []
        text_parts: list[str] = []
        finish_messages: list[str] = []
        candidates = data.get("candidates")
        if not isinstance(candidates, list):
            raise ExternalProviderRequestError("Gemini response payload missing candidates")
        for candidate in candidates:
            if not isinstance(candidate, dict):
                continue
            finish_message = candidate.get("finishMessage")
            finish_reason = candidate.get("finishReason")
            if isinstance(finish_message, str):
                finish_messages.append(finish_message)
            elif isinstance(finish_reason, str):
                finish_messages.append(f"Finish reason: {finish_reason}")
            for part in _iter_response_parts(candidate):
                inline_data = part.get("inline_data") or part.get("inlineData")
                if isinstance(inline_data, dict):
                    encoded = inline_data.get("data")
                    if encoded:
                        image = decode_image_base64(encoded)
                        images.append(ExternalGeneratedImage(image=image, seed=request.seed))
                        continue
                file_data = part.get("fileData") or part.get("file_data")
                if isinstance(file_data, dict):
                    file_uri = file_data.get("fileUri") or file_data.get("file_uri")
                    if isinstance(file_uri, str) and file_uri:
                        raise ExternalProviderRequestError(
                            f"Gemini returned fileUri instead of inline image data: {file_uri}"
                        )
                text = part.get("text")
                if isinstance(text, str):
                    text_parts.append(text)

        if not images:
            self._logger.error("Gemini response contained no images: %s", data)
            detail = ""
            if finish_messages:
                combined = " ".join(message.strip() for message in finish_messages if message.strip())
                if combined:
                    detail = f" Response status: {combined[:500]}"
            elif text_parts:
                combined = " ".join(text_parts).strip()
                if combined:
                    detail = f" Response text: {combined[:500]}"
            raise ExternalProviderRequestError(f"Gemini response contained no images.{detail}")

        return ExternalGenerationResult(
            images=images,
            seed_used=request.seed,
            provider_metadata={"model": request.model.provider_model_id},
        )


def _iter_response_parts(candidate: dict[str, object]) -> list[dict[str, object]]:
    content = candidate.get("content")
    if isinstance(content, dict):
        content_parts = content.get("parts")
        if isinstance(content_parts, list):
            return [part for part in content_parts if isinstance(part, dict)]
    contents = candidate.get("contents")
    if isinstance(contents, list):
        parts: list[dict[str, object]] = []
        for item in contents:
            if not isinstance(item, dict):
                continue
            item_parts = item.get("parts")
            if isinstance(item_parts, list):
                parts.extend([part for part in item_parts if isinstance(part, dict)])
        if parts:
            return parts
    return []


def _select_aspect_ratio(width: int, height: int, allowed: list[str] | None) -> str | None:
    if width <= 0 or height <= 0:
        return None
    ratio = width / height
    default_ratio = _format_aspect_ratio(width, height)
    if not allowed:
        return default_ratio
    parsed = [(value, _parse_ratio(value)) for value in allowed]
    filtered = [(value, parsed_ratio) for value, parsed_ratio in parsed if parsed_ratio is not None]
    if not filtered:
        return default_ratio
    return min(filtered, key=lambda item: abs(item[1] - ratio))[0]


def _format_aspect_ratio(width: int, height: int) -> str | None:
    if width <= 0 or height <= 0:
        return None
    divisor = _gcd(width, height)
    return f"{width // divisor}:{height // divisor}"


def _parse_ratio(value: str) -> float | None:
    if ":" not in value:
        return None
    left, right = value.split(":", 1)
    try:
        numerator = float(left)
        denominator = float(right)
    except ValueError:
        return None
    if denominator == 0:
        return None
    return numerator / denominator


def _parse_retry_after(value: str | None) -> float | None:
    if not value:
        return None
    try:
        return float(value)
    except ValueError:
        return None


def _gcd(a: int, b: int) -> int:
    while b:
        a, b = b, a % b
    return a