feat(skills): add image-gen Gemini API connector script
Python wrapper for the Gemini image generation API, used by the image-gen skill to produce concept art for the 3D pipeline. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Claude Opus 4.6
+248
@@ -0,0 +1,248 @@
|
||||
#!/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 base64
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
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()
|
||||
Reference in New Issue
Block a user