wehub-resource-sync
236 lines
8.8 KiB
Python
236 lines
8.8 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Gemini Image Generation Backend
|
|
|
|
Generates images via the Google GenAI API (Gemini).
|
|
Used by image_gen.py as a backend module.
|
|
|
|
Configuration keys:
|
|
GEMINI_API_KEY (required)
|
|
GEMINI_BASE_URL (optional) Custom API endpoint for proxy services
|
|
GEMINI_MODEL (optional) Override default model
|
|
|
|
Dependencies:
|
|
pip install google-genai Pillow
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
_SCRIPTS_DIR = Path(__file__).resolve().parents[1]
|
|
if str(_SCRIPTS_DIR) not in sys.path:
|
|
sys.path.insert(0, str(_SCRIPTS_DIR))
|
|
|
|
from console_encoding import configure_utf8_stdio # noqa: E402
|
|
|
|
configure_utf8_stdio()
|
|
|
|
if __name__ == "__main__":
|
|
print(__doc__)
|
|
print("Use via: python3 skills/ppt-master/scripts/image_gen.py \"prompt\" --backend gemini")
|
|
raise SystemExit(0 if any(arg in {"-h", "--help", "help"} for arg in sys.argv[1:]) else 1)
|
|
|
|
import os
|
|
import time
|
|
import threading
|
|
from google import genai
|
|
from google.genai import types
|
|
from image_backends.backend_common import (
|
|
MAX_RETRIES,
|
|
is_rate_limit_error,
|
|
normalize_image_size,
|
|
resolve_output_path,
|
|
retry_delay,
|
|
save_image_bytes,
|
|
)
|
|
|
|
|
|
# ╔══════════════════════════════���═══════════════════════════════════╗
|
|
# ║ Constants ║
|
|
# ╚═══════════��═════════════════════════════��════════════════════════╝
|
|
|
|
VALID_ASPECT_RATIOS = [
|
|
"1:1", "1:4", "1:8",
|
|
"2:3", "3:2", "3:4", "4:1", "4:3",
|
|
"4:5", "5:4", "8:1", "9:16", "16:9", "21:9"
|
|
]
|
|
|
|
VALID_IMAGE_SIZES = ["512px", "1K", "2K", "4K"]
|
|
|
|
DEFAULT_MODEL = "gemini-3.1-flash-image-preview"
|
|
|
|
|
|
# ╔══════���═══════════════════════════════════════════════��═══════════╗
|
|
# ║ Image Generation ║
|
|
# ╚══════════════════════════════════════════════════════════════════╝
|
|
|
|
def _generate_image(api_key: str, prompt: str,
|
|
aspect_ratio: str = "1:1", image_size: str = "1K",
|
|
output_dir: str = None, filename: str = None,
|
|
model: str = DEFAULT_MODEL, base_url: str = None) -> str:
|
|
"""
|
|
Image generation via Gemini API (streaming).
|
|
|
|
Returns:
|
|
Path of the saved image file
|
|
|
|
Raises:
|
|
RuntimeError: When generation fails
|
|
"""
|
|
if base_url:
|
|
client = genai.Client(api_key=api_key, http_options={'base_url': base_url})
|
|
else:
|
|
client = genai.Client(api_key=api_key)
|
|
|
|
config_kwargs = {
|
|
"response_modalities": ["IMAGE"],
|
|
"image_config": types.ImageConfig(
|
|
aspect_ratio=aspect_ratio,
|
|
image_size=image_size,
|
|
),
|
|
}
|
|
if "flash" in model.lower():
|
|
config_kwargs["thinking_config"] = types.ThinkingConfig(
|
|
thinking_level="MINIMAL",
|
|
)
|
|
config = types.GenerateContentConfig(**config_kwargs)
|
|
|
|
mode_label = "Proxy Mode" if base_url else "Official Mode"
|
|
print(f"[Gemini - {mode_label}]")
|
|
if base_url:
|
|
print(f" Base URL: {base_url}")
|
|
print(f" Model: {model}")
|
|
print(f" Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}")
|
|
print(f" Aspect Ratio: {aspect_ratio}")
|
|
print(f" Image Size: {image_size}")
|
|
print()
|
|
|
|
start_time = time.time()
|
|
print(f" [..] Generating...", end="", flush=True)
|
|
|
|
heartbeat_stop = threading.Event()
|
|
|
|
def _heartbeat():
|
|
while not heartbeat_stop.is_set():
|
|
heartbeat_stop.wait(5)
|
|
if not heartbeat_stop.is_set():
|
|
elapsed = time.time() - start_time
|
|
print(f" {elapsed:.0f}s...", end="", flush=True)
|
|
|
|
hb_thread = threading.Thread(target=_heartbeat, daemon=True)
|
|
hb_thread.start()
|
|
|
|
last_image_data = None
|
|
chunk_count = 0
|
|
total_bytes = 0
|
|
|
|
for chunk in client.models.generate_content_stream(
|
|
model=model,
|
|
contents=[prompt],
|
|
config=config,
|
|
):
|
|
elapsed = time.time() - start_time
|
|
|
|
if chunk.parts is None:
|
|
continue
|
|
|
|
for part in chunk.parts:
|
|
if part.text is not None:
|
|
print(f"\n Model says: {part.text}", end="", flush=True)
|
|
elif part.inline_data is not None:
|
|
chunk_count += 1
|
|
data_size = len(part.inline_data.data) if part.inline_data.data else 0
|
|
total_bytes += data_size
|
|
size_str = f"{data_size / 1024:.0f}KB" if data_size < 1048576 else f"{data_size / 1048576:.1f}MB"
|
|
print(f"\n [OK] Chunk #{chunk_count} received ({size_str}, {elapsed:.1f}s)", end="", flush=True)
|
|
last_image_data = part
|
|
|
|
heartbeat_stop.set()
|
|
hb_thread.join(timeout=1)
|
|
|
|
elapsed = time.time() - start_time
|
|
print(f"\n [DONE] Stream complete ({elapsed:.1f}s, {chunk_count} chunk(s), {total_bytes / 1024:.0f}KB total)")
|
|
|
|
if last_image_data is not None and last_image_data.inline_data is not None:
|
|
if chunk_count > 1:
|
|
print(f" Keeping the final chunk (highest quality).")
|
|
path = resolve_output_path(prompt, output_dir, filename, ".png")
|
|
return save_image_bytes(
|
|
last_image_data.inline_data.data,
|
|
path,
|
|
content_type=getattr(last_image_data.inline_data, "mime_type", None),
|
|
)
|
|
|
|
raise RuntimeError("No image was generated. The server may have refused the request.")
|
|
|
|
|
|
# ╔���═══════════════════════════════════���═════════════════════════════╗
|
|
# ║ Public Entry Point ║
|
|
# ╚═════════════���═══════════════════════════���════════════════════════╝
|
|
|
|
def generate(prompt: str,
|
|
aspect_ratio: str = "1:1", image_size: str = "1K",
|
|
output_dir: str = None, filename: str = None,
|
|
model: str = None, max_retries: int = MAX_RETRIES) -> str:
|
|
"""
|
|
Gemini image generation with automatic retry.
|
|
|
|
Reads credentials from the current process environment or a `.env` file:
|
|
GEMINI_API_KEY
|
|
GEMINI_BASE_URL
|
|
GEMINI_MODEL (optional override)
|
|
|
|
Args:
|
|
prompt: Prompt text
|
|
aspect_ratio: Aspect ratio (e.g. "16:9", "1:1")
|
|
image_size: Image size ("512px", "1K", "2K", "4K", case-insensitive)
|
|
output_dir: Output directory
|
|
filename: Output filename (without extension)
|
|
model: Model name (default: gemini-3.1-flash-image-preview)
|
|
max_retries: Maximum number of retries
|
|
|
|
Returns:
|
|
Path of the saved image file
|
|
"""
|
|
api_key = os.environ.get("GEMINI_API_KEY")
|
|
base_url = os.environ.get("GEMINI_BASE_URL")
|
|
|
|
if not api_key:
|
|
raise ValueError(
|
|
"No API key found. Set GEMINI_API_KEY in the current environment or a .env file."
|
|
)
|
|
|
|
if model is None:
|
|
model = os.environ.get("GEMINI_MODEL") or DEFAULT_MODEL
|
|
|
|
image_size = normalize_image_size(image_size)
|
|
|
|
if aspect_ratio not in VALID_ASPECT_RATIOS:
|
|
raise ValueError(f"Invalid aspect ratio '{aspect_ratio}'. Valid: {VALID_ASPECT_RATIOS}")
|
|
|
|
if image_size not in VALID_IMAGE_SIZES:
|
|
raise ValueError(f"Invalid image size '{image_size}'. Valid: {VALID_IMAGE_SIZES}")
|
|
|
|
last_error = None
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
return _generate_image(api_key, prompt,
|
|
aspect_ratio, image_size, output_dir,
|
|
filename, model, base_url)
|
|
except Exception as e:
|
|
last_error = e
|
|
if attempt < max_retries and is_rate_limit_error(e):
|
|
delay = retry_delay(attempt, rate_limited=True)
|
|
print(f"\n [WARN] Rate limit hit (attempt {attempt + 1}/{max_retries + 1}). "
|
|
f"Waiting {delay}s before retry...")
|
|
time.sleep(delay)
|
|
elif attempt < max_retries:
|
|
delay = retry_delay(attempt, rate_limited=False)
|
|
print(f"\n [WARN] Error (attempt {attempt + 1}/{max_retries + 1}): {e}. "
|
|
f"Retrying in {delay}s...")
|
|
time.sleep(delay)
|
|
else:
|
|
break
|
|
|
|
raise RuntimeError(f"Failed after {max_retries + 1} attempts. Last error: {last_error}")
|