if@impactframes.ai
7a1f12efaf
Resolved an issue where Gemini API would not return images in 'generate' mode when an input image was present. The logic now only passes images in 'edit' or 'variation' modes. Fixed a bug where the ResolutionSelector's 'size' output was being overridden by a default 'aspect_ratio' in the Gemini config node. Updated presets and example workflows to align with these changes.
412 lines
17 KiB
Python
412 lines
17 KiB
Python
"""Gemini Image Generation API
|
|
Supports both native Gemini multimodal image generation (Gemini 2.0 Flash)
|
|
and dedicated Imagen models (Imagen 3, Imagen 4, Imagen 4 Ultra).
|
|
"""
|
|
|
|
import asyncio
|
|
import base64
|
|
import logging
|
|
from typing import Optional, Dict, Any, List, Union
|
|
from io import BytesIO
|
|
from PIL import Image
|
|
from google import genai
|
|
from google.genai import types
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Aspect ratio to dimensions mapping
|
|
ASPECT_RATIO_DIMENSIONS = {
|
|
"1:1": (1024, 1024),
|
|
"16:9": (1408, 768),
|
|
"9:16": (768, 1408),
|
|
"4:3": (1280, 896),
|
|
"3:4": (896, 1280),
|
|
}
|
|
|
|
def _convert_size_to_aspect_ratio(size: str) -> str:
|
|
"""Convert OpenAI-style size to aspect ratio."""
|
|
size_to_aspect = {
|
|
"1024x1024": "1:1",
|
|
"1792x1024": "16:9",
|
|
"1024x1792": "9:16",
|
|
}
|
|
return size_to_aspect.get(size, "1:1")
|
|
|
|
async def send_gemini_native_image_request(
|
|
api_key: str,
|
|
prompt: str,
|
|
model: str = "gemini-2.5-flash-image-preview", # Updated to latest model
|
|
n: int = 1,
|
|
size: Optional[str] = None,
|
|
aspect_ratio: Optional[str] = None,
|
|
seed: Optional[int] = None,
|
|
temperature: float = 0.7,
|
|
max_tokens: int = 8192,
|
|
input_images_base64: Optional[List[str]] = None, # NEW: Support for input images
|
|
**kwargs
|
|
) -> Dict[str, Any]:
|
|
"""Generate images using Gemini native multimodal capabilities"""
|
|
|
|
if not api_key:
|
|
raise ValueError("Gemini image generation requires a valid API key")
|
|
|
|
# Create client with API key
|
|
client = genai.Client(api_key=api_key)
|
|
|
|
# Determine aspect ratio and dimensions
|
|
if not aspect_ratio and size:
|
|
aspect_ratio = _convert_size_to_aspect_ratio(size)
|
|
elif not aspect_ratio:
|
|
aspect_ratio = "1:1"
|
|
|
|
# Get target dimensions
|
|
target_width, target_height = ASPECT_RATIO_DIMENSIONS.get(aspect_ratio, (1024, 1024))
|
|
logger.info(f"Using resolution {target_width}x{target_height} for aspect ratio {aspect_ratio}")
|
|
|
|
# Build content list with prompt and optional input images
|
|
contents = []
|
|
|
|
# Add input images first if provided (for image editing/variations)
|
|
if input_images_base64:
|
|
from PIL import Image
|
|
import base64
|
|
from io import BytesIO
|
|
|
|
for img_b64 in input_images_base64:
|
|
try:
|
|
# Convert base64 to PIL Image (Gemini SDK format)
|
|
image_bytes = base64.b64decode(img_b64)
|
|
pil_image = Image.open(BytesIO(image_bytes))
|
|
contents.append(pil_image)
|
|
except Exception as e:
|
|
logger.error(f"Error processing input image: {e}")
|
|
continue
|
|
|
|
# Prepare text prompt with dimensions
|
|
if input_images_base64:
|
|
# For image editing/variation, modify the prompt approach
|
|
prompt_text = f"Based on the provided image(s), {prompt}. Generate a high-quality {target_width}x{target_height} image."
|
|
else:
|
|
# For text-to-image generation
|
|
prompt_text = f"Generate a detailed, high-quality image with dimensions {target_width}x{target_height} of: {prompt}"
|
|
|
|
contents.append(prompt_text)
|
|
all_images = []
|
|
|
|
for i in range(n):
|
|
try:
|
|
# Set up generation config
|
|
config_args = {
|
|
"temperature": temperature,
|
|
"response_modalities": ["IMAGE", "TEXT"], # Request both image and text
|
|
"max_output_tokens": max_tokens,
|
|
}
|
|
|
|
# Generate a unique seed for each image if seed is provided
|
|
if seed is not None:
|
|
config_args["seed"] = (seed + i) % (2**31 - 1)
|
|
logger.info(f"Generating image {i+1}/{n} with seed {config_args['seed']}")
|
|
|
|
generation_config = types.GenerateContentConfig(**config_args)
|
|
|
|
# Generate content with images and prompt
|
|
response = client.models.generate_content(
|
|
model=model,
|
|
contents=contents,
|
|
config=generation_config
|
|
)
|
|
|
|
# Debug logging
|
|
logger.info(f"Response type: {type(response)}")
|
|
if hasattr(response, 'candidates'):
|
|
logger.info(f"Number of candidates: {len(response.candidates) if response.candidates else 0}")
|
|
|
|
# Extract images from response
|
|
image_found = False
|
|
if hasattr(response, 'candidates') and response.candidates:
|
|
for candidate in response.candidates:
|
|
if hasattr(candidate, 'content') and hasattr(candidate.content, 'parts'):
|
|
logger.info(f"Number of parts in candidate: {len(candidate.content.parts)}")
|
|
for idx, part in enumerate(candidate.content.parts):
|
|
logger.info(f"Part {idx} type: {type(part)}, has inline_data: {hasattr(part, 'inline_data')}")
|
|
if hasattr(part, 'text') and part.text:
|
|
logger.info(f"Part {idx} contains text: {part.text[:100] if part.text else 'None'}...")
|
|
# Extract image data
|
|
if hasattr(part, 'inline_data') and part.inline_data:
|
|
image_found = True
|
|
try:
|
|
# The image data is already in bytes
|
|
image_bytes = part.inline_data.data
|
|
|
|
# Validate image data
|
|
if not image_bytes or len(image_bytes) == 0:
|
|
logger.warning("Received empty image data from Gemini")
|
|
continue
|
|
|
|
# Convert to base64
|
|
image_base64 = base64.b64encode(image_bytes).decode('utf-8')
|
|
|
|
# Validate base64 encoding
|
|
if image_base64 and len(image_base64) > 100: # Basic sanity check
|
|
all_images.append({"b64_json": image_base64})
|
|
logger.info(f"Successfully processed Gemini image (size: {len(image_bytes)} bytes)")
|
|
else:
|
|
logger.warning("Generated invalid or too small base64 image data")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing Gemini image data: {e}")
|
|
|
|
if not image_found:
|
|
logger.warning(f"No image found in response {i+1}. The model may have returned only text.")
|
|
|
|
# Add small delay between requests if generating multiple
|
|
if i < n - 1:
|
|
await asyncio.sleep(0.5)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error generating image {i+1}: {e}")
|
|
continue
|
|
|
|
if not all_images:
|
|
logger.error("No images were generated")
|
|
return {"error": "No images generated", "data": []}
|
|
|
|
return {"data": all_images}
|
|
|
|
|
|
async def send_imagen_image_request(
|
|
api_key: str,
|
|
prompt: str,
|
|
model: str = "imagen-3.0-generate-002",
|
|
n: int = 1,
|
|
aspect_ratio: str = "1:1",
|
|
person_generation: str = "allow_adult",
|
|
safety_filter_level: str = "block_medium_and_above",
|
|
language: Optional[str] = None,
|
|
seed: Optional[int] = None,
|
|
**kwargs
|
|
) -> Dict[str, Any]:
|
|
"""Generate images using Google's Imagen models"""
|
|
|
|
if not api_key:
|
|
raise ValueError("Imagen generation requires a valid API key")
|
|
|
|
# Create client with API key
|
|
client = genai.Client(api_key=api_key)
|
|
|
|
# Note: For Imagen 4 Ultra, only 1 image can be generated at a time
|
|
if "ultra" in model.lower() and n > 1:
|
|
logger.warning("Imagen 4 Ultra only supports generating 1 image at a time. Setting n=1.")
|
|
n = 1
|
|
|
|
try:
|
|
# Set up generation config for Imagen
|
|
config_args = {
|
|
# Required / supported fields
|
|
"number_of_images": n,
|
|
"aspect_ratio": aspect_ratio,
|
|
}
|
|
|
|
# Optional fields with mapping / validation -------------------------
|
|
# person_generation is allowed unless it's the default (allow_adult)
|
|
if person_generation and person_generation != "allow_adult":
|
|
config_args["person_generation"] = person_generation
|
|
|
|
# Map legacy safety level names (block_few / some / most) ➜ new names
|
|
legacy_map = {
|
|
"block_few": "block_low_and_above",
|
|
"block_some": "block_medium_and_above",
|
|
"block_most": "block_high_and_above",
|
|
}
|
|
if safety_filter_level:
|
|
mapped_level = legacy_map.get(safety_filter_level, safety_filter_level)
|
|
# Imagen preview models currently only support block_low_and_above
|
|
if mapped_level not in {
|
|
"block_low_and_above",
|
|
"block_medium_and_above",
|
|
"block_high_and_above",
|
|
}:
|
|
mapped_level = "block_low_and_above"
|
|
|
|
# Preview Imagen endpoints currently support only block_low_and_above
|
|
if "preview" in model.lower():
|
|
mapped_level = "block_low_and_above"
|
|
config_args["safety_filter_level"] = mapped_level
|
|
|
|
# Language (omit if 'auto')
|
|
if language and language != "auto":
|
|
config_args["language"] = language
|
|
|
|
generation_config = types.GenerateImagesConfig(**config_args)
|
|
|
|
# Generate images using the Imagen-specific method
|
|
response = client.models.generate_images(
|
|
model=model,
|
|
prompt=prompt,
|
|
config=generation_config
|
|
)
|
|
|
|
# Process generated images
|
|
all_images: List[Dict[str, str]] = []
|
|
|
|
if getattr(response, "generated_images", None):
|
|
gen_images = response.generated_images or []
|
|
for generated_image in gen_images:
|
|
try:
|
|
img_obj = getattr(generated_image, "image", None)
|
|
if img_obj is None:
|
|
continue
|
|
|
|
image_bytes = None
|
|
# The SDK may return either raw bytes or a PIL Image
|
|
if isinstance(img_obj, bytes):
|
|
image_bytes = img_obj
|
|
else:
|
|
# It's an Image object. Its `save` method expects a file path.
|
|
import tempfile
|
|
import os
|
|
|
|
fp, temp_path = tempfile.mkstemp(suffix=".png")
|
|
os.close(fp)
|
|
|
|
try:
|
|
# Save to the temporary path
|
|
img_obj.save(temp_path)
|
|
# Read the bytes back
|
|
with open(temp_path, 'rb') as f:
|
|
image_bytes = f.read()
|
|
except Exception as e:
|
|
logger.error(f"Failed to process image via temporary file: {e}")
|
|
finally:
|
|
# Clean up the temp file
|
|
if os.path.exists(temp_path):
|
|
os.unlink(temp_path)
|
|
|
|
if image_bytes:
|
|
image_base64 = base64.b64encode(image_bytes).decode("utf-8")
|
|
all_images.append({"b64_json": image_base64})
|
|
else:
|
|
logger.warning("Could not get image bytes for one returned image.")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing Imagen response item: {e}")
|
|
else:
|
|
# Fallback: some Imagen responses embed images in candidates/parts (rare)
|
|
parts = getattr(response, "candidates", None)
|
|
if parts:
|
|
for cand in parts:
|
|
if hasattr(cand, "content") and hasattr(cand.content, "parts"):
|
|
for part in cand.content.parts:
|
|
if hasattr(part, "inline_data") and part.inline_data:
|
|
try:
|
|
image_bytes = part.inline_data.data
|
|
image_base64 = base64.b64encode(image_bytes).decode("utf-8")
|
|
all_images.append({"b64_json": image_base64})
|
|
except Exception:
|
|
pass
|
|
|
|
if not all_images:
|
|
logger.error("No images were generated by Imagen")
|
|
return {"error": "No images generated", "data": []}
|
|
|
|
return {"data": all_images}
|
|
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
logger.error(f"Imagen generation error: {error_msg}")
|
|
|
|
# Check for specific error types
|
|
if "404" in error_msg or "not found" in error_msg.lower():
|
|
logger.info(f"Model {model} not found, trying alternative names...")
|
|
# Try alternative model names
|
|
alt_models = {
|
|
"imagen-3.0-generate-001": "imagen-3.0-generate-002",
|
|
"imagen-4.0-generate-001": "imagen-4.0-generate-preview-06-06",
|
|
"imagen-4.0-ultra-generate-001": "imagen-4.0-ultra-generate-preview-06-06",
|
|
}
|
|
if model in alt_models:
|
|
return await send_imagen_image_request(
|
|
api_key=api_key,
|
|
prompt=prompt,
|
|
model=alt_models[model],
|
|
n=n,
|
|
aspect_ratio=aspect_ratio,
|
|
person_generation=person_generation,
|
|
safety_filter_level=safety_filter_level,
|
|
language=language,
|
|
seed=seed,
|
|
**kwargs
|
|
)
|
|
|
|
raise
|
|
|
|
|
|
async def send_gemini_image_generation_unified(
|
|
api_key: str,
|
|
model: str,
|
|
prompt: str,
|
|
n: int = 1,
|
|
size: Optional[str] = None,
|
|
aspect_ratio: Optional[str] = None,
|
|
seed: Optional[int] = None,
|
|
edit_mode: bool = False,
|
|
variation_mode: bool = False,
|
|
input_image_base64: Optional[Union[str, List[str]]] = None,
|
|
mask_base64: Optional[str] = None,
|
|
**kwargs
|
|
) -> Dict[str, Any]:
|
|
"""Unified entry point for Gemini/Imagen image generation"""
|
|
|
|
# Normalize model name (remove preview suffixes for routing)
|
|
model_lower = model.lower()
|
|
|
|
# Map preview models to their stable versions if needed
|
|
model_mapping = {
|
|
"imagen-3.0-generate-preview-06-06": "imagen-3.0-generate-002",
|
|
"imagen-3-light-alpha": "imagen-3.0-generate-002", # Legacy naming
|
|
"gemini-2.0-flash-preview-image-generation": "gemini-2.5-flash-image-preview", # Upgrade to latest
|
|
}
|
|
|
|
if model in model_mapping:
|
|
logger.info(f"Mapping model {model} to {model_mapping[model]}")
|
|
model = model_mapping[model]
|
|
|
|
# Route to appropriate function
|
|
if "imagen" in model_lower:
|
|
# Imagen models
|
|
return await send_imagen_image_request(
|
|
api_key=api_key,
|
|
prompt=prompt,
|
|
model=model,
|
|
n=n,
|
|
aspect_ratio=aspect_ratio or _convert_size_to_aspect_ratio(size or "1024x1024"),
|
|
person_generation=kwargs.get("person_generation", "allow_adult"),
|
|
safety_filter_level=kwargs.get("safety_filter_level", "block_medium_and_above"),
|
|
language=kwargs.get("language"),
|
|
seed=seed,
|
|
)
|
|
else:
|
|
# Gemini native models (including image generation models)
|
|
# Support for both gemini-2.5-flash-image-preview and gemini-2.0-flash-preview-image-generation
|
|
|
|
# Prepare input images for native Gemini (only for edit/variation)
|
|
input_images_for_native = None
|
|
if (edit_mode or variation_mode) and input_image_base64:
|
|
if isinstance(input_image_base64, list):
|
|
input_images_for_native = input_image_base64
|
|
else:
|
|
input_images_for_native = [input_image_base64]
|
|
logger.info(f"Gemini native: Using {len(input_images_for_native)} input images for edit/variation.")
|
|
|
|
return await send_gemini_native_image_request(
|
|
api_key=api_key,
|
|
prompt=prompt,
|
|
model=model,
|
|
n=n,
|
|
size=size,
|
|
aspect_ratio=aspect_ratio,
|
|
seed=seed,
|
|
temperature=kwargs.get("temperature", 0.7),
|
|
max_tokens=kwargs.get("max_tokens", 8192),
|
|
input_images_base64=input_images_for_native,
|
|
) |