#!/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()