""" Gemini API Client ================= Client for Google Gemini Image APIs (Flash and Pro models). Handles API communication and response parsing. """ import base64 import logging from io import BytesIO from typing import Optional from PIL import Image from google import genai from google.genai import types from .models import GenerationRequest, GenerationResult logger = logging.getLogger(__name__) class GeminiClient: """ Client for Gemini Image APIs. Supports: - Gemini 2.5 Flash Image (up to ~3 reference images) - Gemini 3 Pro Image Preview (up to 14 reference images, 1K/2K/4K) """ # Model names (updated January 2026) # See: https://ai.google.dev/gemini-api/docs/image-generation MODEL_FLASH = "gemini-2.5-flash-image" # Fast, efficient image generation MODEL_PRO = "gemini-3-pro-image-preview" # Pro quality, advanced text rendering # Valid resolutions for Pro model VALID_RESOLUTIONS = ["1K", "2K", "4K"] # Aspect ratio to dimensions mapping ASPECT_RATIOS = { "1:1": (1024, 1024), "16:9": (1344, 768), "9:16": (768, 1344), "21:9": (1536, 640), # Cinematic ultra-wide "3:2": (1248, 832), "2:3": (832, 1248), "3:4": (864, 1184), "4:3": (1344, 1008), "4:5": (1024, 1280), "5:4": (1280, 1024), } def __init__(self, api_key: str, use_pro_model: bool = False): """ Initialize Gemini client. Args: api_key: Google Gemini API key use_pro_model: If True, use Pro model with enhanced capabilities """ if not api_key: raise ValueError("API key is required for Gemini client") self.api_key = api_key self.use_pro_model = use_pro_model self.client = genai.Client(api_key=api_key) model_name = self.MODEL_PRO if use_pro_model else self.MODEL_FLASH logger.info(f"GeminiClient initialized with model: {model_name}") def generate( self, request: GenerationRequest, resolution: str = "1K" ) -> GenerationResult: """ Generate image using Gemini API. Args: request: GenerationRequest object resolution: Resolution for Pro model ("1K", "2K", "4K") Returns: GenerationResult object """ try: model_name = self.MODEL_PRO if self.use_pro_model else self.MODEL_FLASH logger.info(f"Generating with {model_name}: {request.prompt[:100]}...") # Build contents list contents = self._build_contents(request) # Build config config = self._build_config( request, resolution if self.use_pro_model else None ) # Call API response = self.client.models.generate_content( model=model_name, contents=contents, config=config ) # Parse response return self._parse_response(response) except Exception as e: logger.error(f"Gemini generation failed: {e}", exc_info=True) return GenerationResult.error_result(f"Gemini API error: {str(e)}") def _build_contents(self, request: GenerationRequest) -> list: """Build contents list for API request.""" contents = [] # Add input images if present if request.has_input_images: valid_images = [img for img in request.input_images if img is not None] contents.extend(valid_images) # Add prompt contents.append(request.prompt) return contents def _build_config( self, request: GenerationRequest, resolution: Optional[str] = None ) -> types.GenerateContentConfig: """Build generation config for API request.""" # Parse aspect ratio aspect_ratio = request.aspect_ratio if " " in aspect_ratio: aspect_ratio = aspect_ratio.split()[0] # Build image config image_config_kwargs = {"aspect_ratio": aspect_ratio} # Add resolution for Pro model if resolution and self.use_pro_model: if resolution not in self.VALID_RESOLUTIONS: logger.warning(f"Invalid resolution '{resolution}', defaulting to '1K'") resolution = "1K" image_config_kwargs["output_image_resolution"] = resolution logger.info(f"Pro model resolution: {resolution}") config = types.GenerateContentConfig( temperature=request.temperature, response_modalities=["image", "text"], image_config=types.ImageConfig(**image_config_kwargs) ) return config def _parse_response(self, response) -> GenerationResult: """Parse API response and extract image.""" if response is None: return GenerationResult.error_result("No response from API") if not hasattr(response, 'candidates') or not response.candidates: return GenerationResult.error_result("No candidates in response") candidate = response.candidates[0] # Check finish reason if hasattr(candidate, 'finish_reason'): finish_reason = str(candidate.finish_reason) logger.info(f"Finish reason: {finish_reason}") if 'SAFETY' in finish_reason or 'PROHIBITED' in finish_reason: return GenerationResult.error_result( f"Content blocked by safety filters: {finish_reason}" ) # Check for content if not hasattr(candidate, 'content') or candidate.content is None: finish_reason = getattr(candidate, 'finish_reason', 'UNKNOWN') return GenerationResult.error_result( f"No content in response (finish_reason: {finish_reason})" ) # Extract image from parts if hasattr(candidate.content, 'parts') and candidate.content.parts: for part in candidate.content.parts: if hasattr(part, 'inline_data') and part.inline_data: try: image_data = part.inline_data.data # Handle both bytes and base64 string if isinstance(image_data, str): image_data = base64.b64decode(image_data) # Convert to PIL Image image_buffer = BytesIO(image_data) image = Image.open(image_buffer) image.load() logger.info(f"Image generated: {image.size}, {image.mode}") return GenerationResult.success_result( image=image, message="Generated successfully" ) except Exception as e: logger.error(f"Failed to decode image: {e}") return GenerationResult.error_result( f"Image decoding error: {str(e)}" ) return GenerationResult.error_result("No image data in response") def is_healthy(self) -> bool: """Check if API is accessible.""" return self.api_key is not None and len(self.api_key) > 0 @classmethod def get_dimensions(cls, aspect_ratio: str) -> tuple: """Get pixel dimensions for aspect ratio.""" ratio = aspect_ratio.split()[0] if " " in aspect_ratio else aspect_ratio return cls.ASPECT_RATIOS.get(ratio, (1024, 1024))