166 lines
6.0 KiB
Python
166 lines
6.0 KiB
Python
#!/usr/bin/env -S uv run
|
|
# /// script
|
|
# requires-python = ">=3.10"
|
|
# dependencies = [
|
|
# "httpx>=0.27.0",
|
|
# "python-dotenv>=1.0.0",
|
|
# ]
|
|
# ///
|
|
"""
|
|
Generate plan images via OpenRouter (Gemini 3 image-capable models).
|
|
|
|
Drop-in alternative to generate_gpt_image.py — same CLI signature so the
|
|
planf3 workflows need no changes. Uses your existing OPENROUTER_API_KEY
|
|
instead of a paid OpenAI key.
|
|
|
|
Usage:
|
|
python generate_or_image.py "prompt" output.png [options]
|
|
|
|
Examples:
|
|
python generate_or_image.py "A sunset over mountains" sunset.png
|
|
python generate_or_image.py "Wide architecture diagram" wide.png --size 1536x1024 --quality high
|
|
|
|
Environment:
|
|
OPENROUTER_API_KEY - Required (get from https://openrouter.ai/settings/keys)
|
|
"""
|
|
|
|
import argparse
|
|
import base64
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv(Path.cwd() / ".env")
|
|
|
|
# OpenRouter image-capable models. Gemini 3 Flash Image is cheap + fast;
|
|
# Pro Image is higher quality. Pick via --model or OPENROUTER_IMAGE_MODEL env.
|
|
DEFAULT_MODEL = os.environ.get("OPENROUTER_IMAGE_MODEL", "google/gemini-3.1-flash-image")
|
|
API_BASE = "https://openrouter.ai/api/v1/chat/completions"
|
|
|
|
# OpenRouter uses this to attribute usage in their dashboard.
|
|
HTTP_REFERER = os.environ.get("OPENROUTER_REFERER", "https://github.com/disler/planf3")
|
|
APP_TITLE = "planf3"
|
|
|
|
|
|
def parse_size(size: str) -> str:
|
|
"""Validate and normalize the size argument. We pass it through to the model
|
|
via the prompt; OpenRouter image models derive dimensions from the prompt
|
|
context, so we keep the explicit size in the request text as a hint."""
|
|
if size == "auto":
|
|
return "auto"
|
|
if not re.match(r"^\d+x\d+$", size):
|
|
raise ValueError(f"invalid size '{size}' — expected WxH (e.g. 1536x1024) or 'auto'")
|
|
return size
|
|
|
|
|
|
def generate(prompt: str, output_path: str, size: str, quality: str, model: str) -> str:
|
|
api_key = os.environ.get("OPENROUTER_API_KEY")
|
|
if not api_key:
|
|
raise EnvironmentError(
|
|
"OPENROUTER_API_KEY environment variable not set. "
|
|
"Get one from https://openrouter.ai/settings/keys and add to .env"
|
|
)
|
|
|
|
# Compose the image request. Image-capable Gemini models on OpenRouter
|
|
# return an inline base64 image in the message content when asked.
|
|
size_hint = f" Image dimensions: {size}." if size != "auto" else ""
|
|
quality_hint = f" Quality: {quality}." if quality != "auto" else ""
|
|
user_content = f"Generate a single professional, minimal image:{size_hint}{quality_hint} {prompt}"
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
"HTTP-Referer": HTTP_REFERER,
|
|
"X-Title": APP_TITLE,
|
|
}
|
|
payload = {
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": user_content}],
|
|
# request an image back
|
|
"modalities": ["image", "text"],
|
|
# cap output tokens to control cost — image bytes count against this.
|
|
# Gemini image output is ~4 tokens/px so keep this modest.
|
|
"max_tokens": 4096,
|
|
}
|
|
|
|
print(f"[generate_or_image] model={model} size={size} quality={quality}", file=sys.stderr)
|
|
print(f"[generate_or_image] prompt: {prompt[:120]}{'...' if len(prompt)>120 else ''}", file=sys.stderr)
|
|
|
|
with httpx.Client(timeout=180.0) as client:
|
|
resp = client.post(API_BASE, headers=headers, json=payload)
|
|
|
|
if resp.status_code != 200:
|
|
raise RuntimeError(f"OpenRouter API error {resp.status_code}: {resp.text[:400]}")
|
|
|
|
data = resp.json()
|
|
message = data.get("choices", [{}])[0].get("message", {})
|
|
content = message.get("content", "")
|
|
|
|
# OpenRouter returns image-capable model output in two possible shapes:
|
|
# 1. A list of content parts with type "image_url" (data URI)
|
|
# 2. A markdown string like 
|
|
b64_data = None
|
|
mime = "image/png"
|
|
|
|
if isinstance(content, list):
|
|
for part in content:
|
|
if isinstance(part, dict):
|
|
if part.get("type") == "image_url":
|
|
url = part.get("image_url", {}).get("url", "")
|
|
b64_data, mime = _extract_data_uri(url)
|
|
break
|
|
elif isinstance(content, str):
|
|
b64_data, mime = _extract_data_uri(content)
|
|
|
|
if not b64_data:
|
|
raise RuntimeError(
|
|
"No image returned by model. Response message content:\n"
|
|
+ (json.dumps(content)[:500] if content else "(empty)")
|
|
)
|
|
|
|
# Write the decoded bytes
|
|
out = Path(output_path)
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
out.write_bytes(base64.b64decode(b64_data))
|
|
print(f"[generate_or_image] wrote {out} ({out.stat().st_size} bytes, {mime})", file=sys.stderr)
|
|
return str(out)
|
|
|
|
|
|
def _extract_data_uri(text: str):
|
|
"""Pull base64 image data out of a data: URI or a markdown image with a data URI."""
|
|
if not text:
|
|
return None, "image/png"
|
|
m = re.search(r"data:(image/[a-zA-Z+]+);base64,([A-Za-z0-9+/=\s]+)", text)
|
|
if m:
|
|
mime = m.group(1)
|
|
# strip any whitespace the API may have injected
|
|
b64 = re.sub(r"\s+", "", m.group(2))
|
|
return b64, mime
|
|
return None, "image/png"
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="Generate a plan image via OpenRouter.")
|
|
ap.add_argument("prompt", help="Image prompt")
|
|
ap.add_argument("output_path", help="Where to save the PNG")
|
|
ap.add_argument("--size", default="1536x1024", help="WxH or 'auto' (default 1536x1024)")
|
|
ap.add_argument("--quality", default="high", choices=["auto", "low", "medium", "high"], help="Quality hint")
|
|
ap.add_argument("--model", default=DEFAULT_MODEL, help=f"OpenRouter model id (default {DEFAULT_MODEL})")
|
|
args = ap.parse_args()
|
|
|
|
try:
|
|
size = parse_size(args.size)
|
|
generate(args.prompt, args.output_path, size, args.quality, args.model)
|
|
except Exception as e:
|
|
print(f"ERROR: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|