Add Google Imagen and Google Cloud TTS provider tools
Two new provider tools following the BaseTool pattern with auto-discovery: - google_imagen: Imagen 4 image generation via Generative Language REST API - google_tts: Google Cloud TTS with 700+ voices across 50+ languages Both share GOOGLE_API_KEY env var. Selectors auto-discover them — no selector code changes needed. Docs and contract tests updated. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
calesthio
@@ -5,6 +5,10 @@
|
||||
FAL_KEY= # FLUX images, Google Veo video, Kling video, MiniMax video, Recraft images
|
||||
# Get one at https://fal.ai/dashboard/keys
|
||||
|
||||
# --- Google (one key unlocks image gen + TTS) ---
|
||||
GOOGLE_API_KEY= # Google Imagen images, Google Cloud TTS (700+ voices, 50+ languages)
|
||||
# Get one at https://aistudio.google.com/apikey
|
||||
|
||||
# --- Voice ---
|
||||
ELEVENLABS_API_KEY= # TTS narration, music generation, sound effects
|
||||
OPENAI_API_KEY= # OpenAI TTS fallback and DALL-E image generation
|
||||
|
||||
+2
-2
@@ -243,8 +243,8 @@ Three selector tools abstract multi-provider capabilities. **Selectors auto-disc
|
||||
|
||||
| Selector | Routes to | How it discovers |
|
||||
|----------|-----------|-----------------|
|
||||
| `tts_selector` | All tools with `capability="tts"` | `registry.get_by_capability("tts")` |
|
||||
| `image_selector` | All tools with `capability="image_generation"` | `registry.get_by_capability("image_generation")` |
|
||||
| `tts_selector` | All tools with `capability="tts"` (ElevenLabs, Google TTS, OpenAI, Piper) | `registry.get_by_capability("tts")` |
|
||||
| `image_selector` | All tools with `capability="image_generation"` (FLUX, Google Imagen, DALL-E, Recraft, etc.) | `registry.get_by_capability("image_generation")` |
|
||||
| `video_selector` | All tools with `capability="video_generation"` | `registry.get_by_capability("video_generation")` |
|
||||
|
||||
Selectors route based on: user preference > availability > discovery order. They adapt input schemas between providers transparently.
|
||||
|
||||
+1
-1
@@ -46,7 +46,7 @@ Each tool's `agent_skills[]` field bridges Layer 1 → Layer 3. See `skills/INDE
|
||||
- **Instruction-driven stages:** Each stage has a director skill (MD) that teaches the agent HOW
|
||||
- **Pipeline manifests:** Declarative YAML defining stages, skills, tools, review focus, approval gates
|
||||
- **Capability-first tool design:** Each major family should expose a selector tool plus explicit provider tools
|
||||
- Example: `tts_selector` + `elevenlabs_tts` / `openai_tts` / `piper_tts`
|
||||
- Example: `tts_selector` + `elevenlabs_tts` / `google_tts` / `openai_tts` / `piper_tts`
|
||||
- Example: `video_selector` + `heygen_video` / `wan_video` / `hunyuan_video` / `ltx_video_local` / `ltx_video_modal` / `cogvideo_video`
|
||||
- **Style playbooks:** YAML defining visual language, typography, motion, audio, asset generation constraints
|
||||
- **Artifacts are canonical:** `brief`, `script`, `scene_plan`, `asset_manifest`, `edit_decisions`, `render_report`, `publish_log`
|
||||
|
||||
@@ -21,7 +21,7 @@ Most AI video tools give you a single clip from a prompt. OpenMontage gives you
|
||||
Edit your own talking-head footage. Generate a fully animated explainer from scratch. Cut a 2-hour podcast into a dozen social clips. Translate and dub your content into 10 languages. Build a cinematic brand teaser from stock footage and AI-generated scenes. **If a production team can make it, OpenMontage can orchestrate it.**
|
||||
|
||||
- **11 production pipelines** — explainers, talking heads, screen demos, cinematic trailers, animations, podcasts, localization, and more
|
||||
- **47 production tools** — spanning video generation, image creation, text-to-speech, music, audio mixing, subtitles, enhancement, and analysis
|
||||
- **49 production tools** — spanning video generation, image creation, text-to-speech, music, audio mixing, subtitles, enhancement, and analysis
|
||||
- **Live web research built in** — before writing a single word of script, the agent runs 15-25+ web searches across YouTube, Reddit, news sites, and academic sources to ground your video in real, current data
|
||||
- **Both free/local AND cloud providers** — every capability supports open-source local alternatives alongside premium APIs. Use what you have.
|
||||
- **No vendor lock-in** — swap providers freely. The selector pattern auto-routes to whatever's available on your machine.
|
||||
@@ -46,10 +46,11 @@ Edit your own talking-head footage. Generate a fully animated explainer from scr
|
||||
| **Pexels** | Stock | Free stock footage |
|
||||
| **Pixabay** | Stock | Free stock footage |
|
||||
|
||||
### Image Generation (7 providers)
|
||||
### Image Generation (8 providers)
|
||||
| Provider | Type | Notes |
|
||||
|----------|------|-------|
|
||||
| **FLUX** | Cloud API | State-of-the-art quality |
|
||||
| **Google Imagen** | Cloud API | Imagen 4 — high-quality, multiple aspect ratios |
|
||||
| **DALL-E 3** | Cloud API | OpenAI's image model |
|
||||
| **Recraft** | Cloud API | Design-focused generation |
|
||||
| **Local Diffusion** | Local GPU | Stable Diffusion, free |
|
||||
@@ -57,10 +58,11 @@ Edit your own talking-head footage. Generate a fully animated explainer from scr
|
||||
| **Pixabay** | Stock | Free stock images |
|
||||
| **ManimCE** | Local | Mathematical animations |
|
||||
|
||||
### Text-to-Speech (3 providers)
|
||||
### Text-to-Speech (4 providers)
|
||||
| Provider | Type | Notes |
|
||||
|----------|------|-------|
|
||||
| **ElevenLabs** | Cloud API | Premium voice quality |
|
||||
| **Google TTS** | Cloud API | 700+ voices, 50+ languages — best for localization |
|
||||
| **OpenAI TTS** | Cloud API | Fast, affordable |
|
||||
| **Piper** | Local | Completely free, offline |
|
||||
|
||||
@@ -223,6 +225,7 @@ SUNO_API_KEY=your-key # Suno AI — full songs, instrumentals, any genr
|
||||
# Voice & images:
|
||||
ELEVENLABS_API_KEY=your-key # Premium TTS, AI music, sound effects
|
||||
OPENAI_API_KEY=your-key # OpenAI TTS, DALL-E 3 images
|
||||
GOOGLE_API_KEY=your-key # Google Imagen images, Google TTS (700+ voices)
|
||||
|
||||
# More video providers:
|
||||
HEYGEN_API_KEY=your-key # HeyGen — VEO, Sora, Runway, Kling via single gateway
|
||||
@@ -260,10 +263,10 @@ The agent will:
|
||||
|
||||
```
|
||||
OpenMontage/
|
||||
├── tools/ # 46 Python tools (the agent's hands)
|
||||
├── tools/ # 48 Python tools (the agent's hands)
|
||||
│ ├── video/ # 12 video gen providers + compose, stitch, trim
|
||||
│ ├── audio/ # 3 TTS providers + Suno/ElevenLabs music, mixing, enhancement
|
||||
│ ├── graphics/ # 7 image gen providers + diagrams, code snippets, math
|
||||
│ ├── audio/ # 4 TTS providers + Suno/ElevenLabs music, mixing, enhancement
|
||||
│ ├── graphics/ # 8 image gen providers + diagrams, code snippets, math
|
||||
│ ├── enhancement/ # Upscale, bg remove, face enhance, color grade
|
||||
│ ├── analysis/ # Transcription, scene detect, frame sampling
|
||||
│ ├── avatar/ # Talking head, lip sync
|
||||
|
||||
@@ -40,7 +40,7 @@ OpenMontage/
|
||||
│ ├── env_loader.py # .env variable management
|
||||
│ └── providers/ # (Reserved for future provider abstractions)
|
||||
│
|
||||
├── tools/ # 55+ Python tool implementations
|
||||
├── tools/ # 57+ Python tool implementations
|
||||
│ ├── base_tool.py # Abstract base class — the tool contract
|
||||
│ ├── tool_registry.py # Auto-discovery singleton registry
|
||||
│ ├── cost_tracker.py # Budget governance (estimate → reserve → reconcile)
|
||||
@@ -142,8 +142,8 @@ Three selector tools abstract multi-provider capabilities:
|
||||
|
||||
| Selector | Capability | Providers (priority order) |
|
||||
|----------|-----------|---------------------------|
|
||||
| `tts_selector` | Text-to-speech | ElevenLabs > OpenAI > Piper (offline) |
|
||||
| `image_selector` | Image generation | FLUX > DALL-E > Recraft > LocalDiffusion > Pexels/Pixabay (stock) |
|
||||
| `tts_selector` | Text-to-speech | ElevenLabs > Google TTS > OpenAI > Piper (offline) |
|
||||
| `image_selector` | Image generation | FLUX > Google Imagen > DALL-E > Recraft > LocalDiffusion > Pexels/Pixabay (stock) |
|
||||
| `video_selector` | Video generation | Kling > Runway > VEO > MiniMax > HeyGen > LTX (modal) > LTX (local) > CogVideo > Hunyuan > WAN > Pexels/Pixabay (stock) |
|
||||
|
||||
Selectors route based on: user preference > availability > fallback order. They adapt input schemas between providers transparently.
|
||||
@@ -152,13 +152,13 @@ Selectors route based on: user preference > availability > fallback order. They
|
||||
|
||||
**Analysis (4):** transcriber (WhisperX), scene_detect, frame_sampler, video_understand (CLIP/BLIP-2)
|
||||
|
||||
**Audio (7):** elevenlabs_tts, openai_tts, piper_tts, tts_selector, music_gen, audio_mixer, audio_enhance
|
||||
**Audio (8):** elevenlabs_tts, google_tts, openai_tts, piper_tts, tts_selector, music_gen, audio_mixer, audio_enhance
|
||||
|
||||
**Avatar (2):** talking_head (SadTalker/MuseTalk), lip_sync (Wav2Lip)
|
||||
|
||||
**Enhancement (5):** upscale (Real-ESRGAN), bg_remove (rembg/U2Net), face_enhance, face_restore (CodeFormer/GFPGAN), color_grade (FFmpeg LUTs)
|
||||
|
||||
**Graphics (11):** flux_image, openai_image, recraft_image, local_diffusion, pexels_image, pixabay_image, image_selector, code_snippet, diagram_gen, math_animate (ManimCE), image_gen (deprecated)
|
||||
**Graphics (12):** flux_image, google_imagen, openai_image, recraft_image, local_diffusion, pexels_image, pixabay_image, image_selector, code_snippet, diagram_gen, math_animate (ManimCE), image_gen (deprecated)
|
||||
|
||||
**Subtitle (1):** subtitle_gen
|
||||
|
||||
@@ -382,6 +382,7 @@ All config is validated via Pydantic models in `lib/config_model.py`.
|
||||
| `HEYGEN_API_KEY` | heygen_video | Multi-provider video generation |
|
||||
| `PEXELS_API_KEY` | pexels_image, pexels_video | Stock media |
|
||||
| `PIXABAY_API_KEY` | pixabay_image, pixabay_video | Stock media |
|
||||
| `GOOGLE_API_KEY` | google_imagen, google_tts | Google Imagen images, Google Cloud TTS |
|
||||
| `RUNWAY_API_KEY` | runway_video | Runway Gen-4 direct |
|
||||
| `MODAL_LTX2_ENDPOINT_URL` | ltx_video_modal | Self-hosted LTX-2 |
|
||||
| `VIDEO_GEN_LOCAL_ENABLED` | local video tools | Enable local GPU generation |
|
||||
|
||||
@@ -143,7 +143,7 @@ class TestCapabilityMetadata:
|
||||
catalog = reg.capability_catalog()
|
||||
assert "tts" in catalog
|
||||
providers = {item["provider"] for item in catalog["tts"] if item["provider"] != "selector"}
|
||||
assert providers == {"elevenlabs", "openai", "piper"}
|
||||
assert providers == {"elevenlabs", "google_tts", "openai", "piper"}
|
||||
|
||||
|
||||
# ---- Animated Explainer Pipeline ----
|
||||
|
||||
@@ -0,0 +1,223 @@
|
||||
"""Google Cloud Text-to-Speech provider tool.
|
||||
|
||||
Google TTS offers 700+ voices across 50+ languages, including Standard,
|
||||
WaveNet, Neural2, Studio, and Journey voice types — strong for localization.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class GoogleTTS(BaseTool):
|
||||
name = "google_tts"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.VOICE
|
||||
capability = "tts"
|
||||
provider = "google_tts"
|
||||
stability = ToolStability.BETA
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set GOOGLE_API_KEY to your Google Cloud API key with Text-to-Speech enabled.\n"
|
||||
" Enable the API at https://console.cloud.google.com/apis/library/texttospeech.googleapis.com\n"
|
||||
" Or use GOOGLE_APPLICATION_CREDENTIALS for service account auth."
|
||||
)
|
||||
fallback = "openai_tts"
|
||||
fallback_tools = ["openai_tts", "elevenlabs_tts", "piper_tts"]
|
||||
agent_skills = ["text-to-speech"]
|
||||
|
||||
capabilities = [
|
||||
"text_to_speech",
|
||||
"voice_selection",
|
||||
"ssml_support",
|
||||
"multilingual",
|
||||
]
|
||||
supports = {
|
||||
"voice_cloning": False,
|
||||
"multilingual": True,
|
||||
"offline": False,
|
||||
"native_audio": True,
|
||||
"ssml": True,
|
||||
}
|
||||
best_for = [
|
||||
"localization — 700+ voices across 50+ languages",
|
||||
"affordable high-quality TTS (Neural2, WaveNet)",
|
||||
"Google ecosystem integration",
|
||||
]
|
||||
not_good_for = [
|
||||
"voice cloning",
|
||||
"fully offline production",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["text"],
|
||||
"properties": {
|
||||
"text": {"type": "string", "description": "Text to convert to speech"},
|
||||
"voice": {
|
||||
"type": "string",
|
||||
"default": "en-US-Neural2-D",
|
||||
"description": "Voice name (e.g. en-US-Neural2-D, en-US-Studio-O, en-GB-WaveNet-A)",
|
||||
},
|
||||
"language_code": {
|
||||
"type": "string",
|
||||
"default": "en-US",
|
||||
"description": "BCP-47 language code (e.g. en-US, es-ES, ja-JP, fr-FR)",
|
||||
},
|
||||
"speaking_rate": {
|
||||
"type": "number",
|
||||
"default": 1.0,
|
||||
"minimum": 0.25,
|
||||
"maximum": 4.0,
|
||||
"description": "Speaking speed. 1.0 = normal, 0.5 = half speed, 2.0 = double speed",
|
||||
},
|
||||
"pitch": {
|
||||
"type": "number",
|
||||
"default": 0.0,
|
||||
"minimum": -20.0,
|
||||
"maximum": 20.0,
|
||||
"description": "Pitch adjustment in semitones. 0.0 = default",
|
||||
},
|
||||
"audio_encoding": {
|
||||
"type": "string",
|
||||
"default": "MP3",
|
||||
"enum": ["MP3", "LINEAR16", "OGG_OPUS", "MULAW", "ALAW"],
|
||||
"description": "Audio output encoding format",
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=50, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["text", "voice", "language_code", "speaking_rate", "pitch"]
|
||||
side_effects = ["writes audio file to output_path", "calls Google Cloud TTS API"]
|
||||
user_visible_verification = ["Listen to generated audio for natural speech quality"]
|
||||
|
||||
# Extension mapping for audio encodings
|
||||
_EXT_MAP = {
|
||||
"MP3": "mp3",
|
||||
"LINEAR16": "wav",
|
||||
"OGG_OPUS": "ogg",
|
||||
"MULAW": "wav",
|
||||
"ALAW": "wav",
|
||||
}
|
||||
|
||||
def _get_api_key(self) -> str | None:
|
||||
return os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._get_api_key() or os.environ.get("GOOGLE_APPLICATION_CREDENTIALS"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
text = inputs.get("text", "")
|
||||
char_count = len(text)
|
||||
voice = inputs.get("voice", "en-US-Neural2-D")
|
||||
# Pricing per million characters (approximate)
|
||||
if "Studio" in voice:
|
||||
rate_per_char = 0.000160 # $160/1M chars
|
||||
elif "Neural2" in voice or "Journey" in voice:
|
||||
rate_per_char = 0.000016 # $16/1M chars
|
||||
elif "WaveNet" in voice:
|
||||
rate_per_char = 0.000016 # $16/1M chars
|
||||
else:
|
||||
rate_per_char = 0.000004 # $4/1M chars (Standard)
|
||||
return round(char_count * rate_per_char, 4)
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = self._get_api_key()
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="No Google API key found. " + self.install_instructions,
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
try:
|
||||
result = self._generate(inputs, api_key)
|
||||
except Exception as exc:
|
||||
return ToolResult(success=False, error=f"Google TTS failed: {exc}")
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
result.cost_usd = self.estimate_cost(inputs)
|
||||
return result
|
||||
|
||||
def _generate(self, inputs: dict[str, Any], api_key: str) -> ToolResult:
|
||||
import requests
|
||||
|
||||
text = inputs["text"]
|
||||
voice_name = inputs.get("voice", "en-US-Neural2-D")
|
||||
language_code = inputs.get("language_code", "en-US")
|
||||
speaking_rate = inputs.get("speaking_rate", 1.0)
|
||||
pitch = inputs.get("pitch", 0.0)
|
||||
audio_encoding = inputs.get("audio_encoding", "MP3")
|
||||
|
||||
payload = {
|
||||
"input": {"text": text},
|
||||
"voice": {
|
||||
"languageCode": language_code,
|
||||
"name": voice_name,
|
||||
},
|
||||
"audioConfig": {
|
||||
"audioEncoding": audio_encoding,
|
||||
"speakingRate": speaking_rate,
|
||||
"pitch": pitch,
|
||||
},
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
"https://texttospeech.googleapis.com/v1/text:synthesize",
|
||||
headers={"Content-Type": "application/json"},
|
||||
params={"key": api_key},
|
||||
json=payload,
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
audio_content = base64.b64decode(response.json()["audioContent"])
|
||||
|
||||
ext = self._EXT_MAP.get(audio_encoding, "mp3")
|
||||
output_path = Path(inputs.get("output_path", f"tts_output.{ext}"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(audio_content)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": self.provider,
|
||||
"voice": voice_name,
|
||||
"language_code": language_code,
|
||||
"text_length": len(text),
|
||||
"output": str(output_path),
|
||||
"format": audio_encoding,
|
||||
"speaking_rate": speaking_rate,
|
||||
"pitch": pitch,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
model=f"google-tts/{voice_name}",
|
||||
)
|
||||
@@ -0,0 +1,221 @@
|
||||
"""Google Imagen image generation via Gemini API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
# Aspect ratio to approximate pixel dimensions (for cost/reporting only)
|
||||
ASPECT_RATIOS = {
|
||||
"1:1": (1024, 1024),
|
||||
"3:4": (896, 1152),
|
||||
"4:3": (1152, 896),
|
||||
"9:16": (768, 1344),
|
||||
"16:9": (1344, 768),
|
||||
}
|
||||
|
||||
|
||||
def _dims_to_aspect_ratio(width: int, height: int) -> str:
|
||||
"""Convert width/height to the nearest supported aspect ratio."""
|
||||
target = width / height
|
||||
best = "1:1"
|
||||
best_diff = float("inf")
|
||||
for ratio, (w, h) in ASPECT_RATIOS.items():
|
||||
diff = abs(target - w / h)
|
||||
if diff < best_diff:
|
||||
best_diff = diff
|
||||
best = ratio
|
||||
return best
|
||||
|
||||
|
||||
class GoogleImagen(BaseTool):
|
||||
name = "google_imagen"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "image_generation"
|
||||
provider = "google_imagen"
|
||||
stability = ToolStability.BETA
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = [] # checked dynamically via env var
|
||||
install_instructions = (
|
||||
"Set GOOGLE_API_KEY (or GEMINI_API_KEY) to your Google AI API key.\n"
|
||||
" Get one at https://aistudio.google.com/apikey"
|
||||
)
|
||||
agent_skills = []
|
||||
|
||||
capabilities = ["generate_image", "generate_illustration", "text_to_image"]
|
||||
supports = {
|
||||
"negative_prompt": False,
|
||||
"seed": False,
|
||||
"custom_size": False,
|
||||
"aspect_ratio": True,
|
||||
}
|
||||
best_for = [
|
||||
"high-quality photorealistic images",
|
||||
"Google ecosystem integration",
|
||||
"fast generation with multiple aspect ratios",
|
||||
]
|
||||
not_good_for = [
|
||||
"negative prompt control (not supported)",
|
||||
"exact pixel dimensions (uses aspect ratios)",
|
||||
"offline generation",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string", "description": "Image description (max 480 tokens)"},
|
||||
"aspect_ratio": {
|
||||
"type": "string",
|
||||
"enum": ["1:1", "3:4", "4:3", "9:16", "16:9"],
|
||||
"default": "1:1",
|
||||
"description": "Aspect ratio of generated image",
|
||||
},
|
||||
"width": {
|
||||
"type": "integer",
|
||||
"description": "Desired width in pixels — mapped to nearest aspect ratio",
|
||||
},
|
||||
"height": {
|
||||
"type": "integer",
|
||||
"description": "Desired height in pixels — mapped to nearest aspect ratio",
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"imagen-4.0-generate-001",
|
||||
"imagen-4.0-fast-generate-001",
|
||||
"imagen-4.0-ultra-generate-001",
|
||||
],
|
||||
"default": "imagen-4.0-generate-001",
|
||||
"description": "Imagen model variant",
|
||||
},
|
||||
"number_of_images": {
|
||||
"type": "integer",
|
||||
"default": 1,
|
||||
"minimum": 1,
|
||||
"maximum": 4,
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=100, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["prompt", "aspect_ratio", "model"]
|
||||
side_effects = ["writes image file to output_path", "calls Google Generative AI API"]
|
||||
user_visible_verification = ["Inspect generated image for relevance and quality"]
|
||||
|
||||
def _get_api_key(self) -> str | None:
|
||||
return os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._get_api_key():
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
model = inputs.get("model", "imagen-4.0-generate-001")
|
||||
n = inputs.get("number_of_images", 1)
|
||||
if "ultra" in model:
|
||||
return 0.06 * n
|
||||
if "fast" in model:
|
||||
return 0.02 * n
|
||||
return 0.04 * n
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = self._get_api_key()
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="No Google API key found. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
model = inputs.get("model", "imagen-4.0-generate-001")
|
||||
prompt = inputs["prompt"]
|
||||
|
||||
# Resolve aspect ratio: explicit > derived from width/height > default
|
||||
if "aspect_ratio" in inputs:
|
||||
aspect_ratio = inputs["aspect_ratio"]
|
||||
elif "width" in inputs and "height" in inputs:
|
||||
aspect_ratio = _dims_to_aspect_ratio(inputs["width"], inputs["height"])
|
||||
else:
|
||||
aspect_ratio = "1:1"
|
||||
|
||||
number_of_images = inputs.get("number_of_images", 1)
|
||||
|
||||
parameters: dict[str, Any] = {
|
||||
"sampleCount": number_of_images,
|
||||
"aspectRatio": aspect_ratio,
|
||||
}
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"https://generativelanguage.googleapis.com/v1beta/models/{model}:predict",
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"x-goog-api-key": api_key,
|
||||
},
|
||||
json={
|
||||
"instances": [{"prompt": prompt}],
|
||||
"parameters": parameters,
|
||||
},
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
predictions = data.get("predictions", [])
|
||||
if not predictions:
|
||||
return ToolResult(success=False, error="No images returned from Imagen API")
|
||||
|
||||
image_bytes = base64.b64decode(
|
||||
predictions[0]["bytesBase64Encoded"]
|
||||
)
|
||||
|
||||
output_path = Path(inputs.get("output_path", "generated_image.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_bytes)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Imagen generation failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "google_imagen",
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"output": str(output_path),
|
||||
"images_generated": len(predictions),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=self.estimate_cost(inputs),
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
model=model,
|
||||
)
|
||||
Reference in New Issue
Block a user