jpmschweitzerandClaude Opus 4.6
980201fb49
Add venv auto-activation to audio_batch, audio_connector, image_connector, qdrant_connector, and trellis_connector. Scripts re-exec into .venv/bin/python when invoked outside the venv, with a friendly error if .venv is missing. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
257 lines
8.0 KiB
Python
Executable File
257 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.0 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 # noqa: E402
|
|
|
|
ensure_venv()
|
|
|
|
import base64
|
|
import json
|
|
import os
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
CONFIG_PATH = os.path.join(os.path.dirname(__file__), "config.json")
|
|
DEFAULT_OUTPUT_DIR = os.path.expanduser("~/Pictures/mcp-images")
|
|
|
|
|
|
def get_api_key():
|
|
"""Get Gemini API key from env or config."""
|
|
key = os.environ.get("GEMINI_API_KEY")
|
|
if key:
|
|
return key
|
|
try:
|
|
with open(CONFIG_PATH) as f:
|
|
config = json.load(f)
|
|
return config.get("gemini_api_key", "")
|
|
except Exception:
|
|
pass
|
|
print(json.dumps({
|
|
"ok": False,
|
|
"error": "No GEMINI_API_KEY found in environment or config.json"
|
|
}, indent=2))
|
|
sys.exit(1)
|
|
|
|
|
|
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 aspect ratio and size hints
|
|
enhanced_prompt = prompt
|
|
if aspect_ratio and aspect_ratio != "1:1":
|
|
enhanced_prompt += f" Aspect ratio: {aspect_ratio}."
|
|
if image_size:
|
|
enhanced_prompt += f" Resolution: {image_size}."
|
|
|
|
parts.append({"text": enhanced_prompt})
|
|
|
|
payload = json.dumps({
|
|
"contents": [{"parts": parts}],
|
|
"generationConfig": {
|
|
"responseModalities": ["TEXT", "IMAGE"],
|
|
}
|
|
})
|
|
|
|
req = urllib.request.Request(
|
|
url,
|
|
data=payload.encode(),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST"
|
|
)
|
|
|
|
print(f"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()
|