jpmschweitzerandClaude Fable 5
33c14531da
Live-command corrections (atlas corridor-status, real body-ID naming), failure-proofed glb-gen/sprite-gen render scripts, Trellis API reference extracted. image-gen: fixed the output-path bug and de-forked the local image_connector.py to the canonical tooling/db/ copy. ticket skill consolidated to point at ticket-cli.md (setparent-none fix applied there too). Part of T-1099. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
248 lines
8.0 KiB
Python
Executable File
248 lines
8.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Gemini image generator connector — direct API wrapper.
|
|
|
|
Generates images via Google's gemini-2.5-flash-image generation API.
|
|
API key from GEMINI_API_KEY env var or config.json.
|
|
|
|
Usage:
|
|
python3 image_connector.py health
|
|
python3 image_connector.py generate "prompt" [--output file.png] [--aspect 1:1] [--size 1K] [--input image.png]
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
|
|
from common import ensure_venv, get_api_key as _get_api_key # noqa: E402
|
|
|
|
ensure_venv()
|
|
|
|
import base64
|
|
import json
|
|
import os
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
DEFAULT_OUTPUT_DIR = os.path.expanduser("~/Pictures/mcp-images")
|
|
|
|
|
|
def get_api_key() -> str:
|
|
"""Get Gemini API key from env or config."""
|
|
return _get_api_key("GEMINI_API_KEY", "gemini_api_key")
|
|
|
|
|
|
def health():
|
|
"""Check if the Gemini API is reachable with the configured key."""
|
|
key = get_api_key()
|
|
url = f"https://generativelanguage.googleapis.com/v1beta/models?key={key}"
|
|
try:
|
|
req = urllib.request.Request(url, method="GET")
|
|
with urllib.request.urlopen(req, timeout=10) as resp:
|
|
data = json.loads(resp.read())
|
|
models = [m.get("name", "") for m in data.get("models", [])
|
|
if "imagen" in m.get("name", "").lower()
|
|
or "flash" in m.get("name", "").lower()]
|
|
print(json.dumps({
|
|
"ok": True,
|
|
"api": "gemini",
|
|
"image_capable_models": models[:5],
|
|
}, indent=2))
|
|
except Exception as e:
|
|
print(json.dumps({
|
|
"ok": False,
|
|
"error": str(e)
|
|
}, indent=2))
|
|
sys.exit(1)
|
|
|
|
|
|
def generate(prompt, output=None, aspect_ratio="1:1", image_size=None,
|
|
input_image=None):
|
|
"""
|
|
Generate an image from a text prompt using Gemini.
|
|
|
|
Args:
|
|
prompt: Text description of the image to generate
|
|
output: Output file path (default: auto-named in ~/Pictures/mcp-images/)
|
|
aspect_ratio: Aspect ratio (1:1, 16:9, 3:2, etc.)
|
|
image_size: Resolution hint (1K, 2K, 4K) - may not be honored
|
|
input_image: Optional input image path for image-to-image generation
|
|
"""
|
|
key = get_api_key()
|
|
|
|
# Gemini image generation model
|
|
model = "gemini-2.5-flash-image"
|
|
url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={key}"
|
|
|
|
if output is None:
|
|
safe = "".join(c if c.isalnum() or c in "-_ " else "" for c in prompt[:40])
|
|
safe = safe.strip().replace(" ", "_").lower()
|
|
os.makedirs(DEFAULT_OUTPUT_DIR, exist_ok=True)
|
|
output = os.path.join(DEFAULT_OUTPUT_DIR, f"{safe}.png")
|
|
|
|
# Build the request
|
|
parts = []
|
|
|
|
# Add input image if provided (image-to-image)
|
|
if input_image:
|
|
if not os.path.isfile(input_image):
|
|
print(json.dumps({"ok": False, "error": f"Input image not found: {input_image}"}, indent=2))
|
|
sys.exit(1)
|
|
with open(input_image, "rb") as f:
|
|
image_data = base64.b64encode(f.read()).decode("utf-8")
|
|
# Detect mime type
|
|
ext = os.path.splitext(input_image)[1].lower()
|
|
mime = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg",
|
|
".webp": "image/webp"}.get(ext, "image/png")
|
|
parts.append({
|
|
"inlineData": {
|
|
"mimeType": mime,
|
|
"data": image_data
|
|
}
|
|
})
|
|
|
|
# Build enhanced prompt with a size hint. Unlike aspect ratio below, Gemini
|
|
# has no dedicated resolution parameter for this model — this is a
|
|
# best-effort prompt hint only and may not be honored.
|
|
enhanced_prompt = prompt
|
|
if image_size:
|
|
enhanced_prompt += f" Resolution: {image_size}."
|
|
|
|
parts.append({"text": enhanced_prompt})
|
|
|
|
generation_config = {"responseModalities": ["TEXT", "IMAGE"]}
|
|
if aspect_ratio:
|
|
# Real API parameter (not a prompt hint). Valid values: 1:1, 3:2,
|
|
# 2:3, 3:4, 4:3, 4:5, 5:4, 9:16, 16:9, 21:9.
|
|
# https://ai.google.dev/gemini-api/docs/image-generation
|
|
generation_config["imageConfig"] = {"aspectRatio": aspect_ratio}
|
|
|
|
payload = json.dumps({
|
|
"contents": [{"parts": parts}],
|
|
"generationConfig": generation_config
|
|
})
|
|
|
|
req = urllib.request.Request(
|
|
url,
|
|
data=payload.encode(),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST"
|
|
)
|
|
|
|
print("Generating image...", file=sys.stderr)
|
|
print(f" Prompt: {prompt}", file=sys.stderr)
|
|
if input_image:
|
|
print(f" Input image: {input_image}", file=sys.stderr)
|
|
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=120) as resp:
|
|
result = json.loads(resp.read())
|
|
except urllib.error.HTTPError as e:
|
|
body = e.read().decode("utf-8", errors="replace")
|
|
print(json.dumps({
|
|
"ok": False,
|
|
"error": f"API error {e.code}: {e.reason}",
|
|
"details": body[:500]
|
|
}, indent=2))
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
print(json.dumps({"ok": False, "error": str(e)}, indent=2))
|
|
sys.exit(1)
|
|
|
|
# Extract image data from response
|
|
candidates = result.get("candidates", [])
|
|
if not candidates:
|
|
print(json.dumps({
|
|
"ok": False,
|
|
"error": "No candidates in response",
|
|
"response": json.dumps(result)[:500]
|
|
}, indent=2))
|
|
sys.exit(1)
|
|
|
|
image_saved = False
|
|
text_response = ""
|
|
|
|
for candidate in candidates:
|
|
content = candidate.get("content", {})
|
|
for part in content.get("parts", []):
|
|
if "inlineData" in part:
|
|
# Image data
|
|
image_b64 = part["inlineData"]["data"]
|
|
image_bytes = base64.b64decode(image_b64)
|
|
os.makedirs(os.path.dirname(os.path.abspath(output)), exist_ok=True)
|
|
with open(output, "wb") as f:
|
|
f.write(image_bytes)
|
|
image_saved = True
|
|
elif "text" in part:
|
|
text_response += part["text"]
|
|
|
|
if not image_saved:
|
|
print(json.dumps({
|
|
"ok": False,
|
|
"error": "No image data in response",
|
|
"text_response": text_response[:500],
|
|
"response": json.dumps(result)[:500]
|
|
}, indent=2))
|
|
sys.exit(1)
|
|
|
|
file_size = os.path.getsize(output)
|
|
print(json.dumps({
|
|
"ok": True,
|
|
"file": output,
|
|
"size_bytes": file_size,
|
|
"prompt": prompt,
|
|
"aspect_ratio": aspect_ratio,
|
|
}, indent=2))
|
|
|
|
|
|
def main():
|
|
if len(sys.argv) < 2:
|
|
print("Usage:")
|
|
print(" image_connector.py health")
|
|
print(" image_connector.py generate 'prompt' [--output file.png] [--aspect 1:1] [--size 1K] [--input image.png]")
|
|
sys.exit(1)
|
|
|
|
cmd = sys.argv[1]
|
|
|
|
if cmd == "health":
|
|
health()
|
|
elif cmd == "generate":
|
|
if len(sys.argv) < 3:
|
|
print("Error: prompt required", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
prompt = sys.argv[2]
|
|
output = None
|
|
aspect_ratio = "1:1"
|
|
image_size = None
|
|
input_image = None
|
|
|
|
i = 3
|
|
while i < len(sys.argv):
|
|
if sys.argv[i] == "--output" and i + 1 < len(sys.argv):
|
|
output = sys.argv[i + 1]
|
|
i += 2
|
|
elif sys.argv[i] == "--aspect" and i + 1 < len(sys.argv):
|
|
aspect_ratio = sys.argv[i + 1]
|
|
i += 2
|
|
elif sys.argv[i] == "--size" and i + 1 < len(sys.argv):
|
|
image_size = sys.argv[i + 1]
|
|
i += 2
|
|
elif sys.argv[i] == "--input" and i + 1 < len(sys.argv):
|
|
input_image = sys.argv[i + 1]
|
|
i += 2
|
|
else:
|
|
print(f"Unknown argument: {sys.argv[i]}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
generate(prompt, output=output, aspect_ratio=aspect_ratio,
|
|
image_size=image_size, input_image=input_image)
|
|
else:
|
|
print(f"Unknown command: {cmd}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|