Files
omniclaw-skills/skills/sub2api-gpt-image-2/scripts/generate_image.py
2026-04-24 01:12:20 -07:00

124 lines
4.0 KiB
Python
Executable File

#!/usr/bin/env python3
import argparse
import base64
import json
import os
import sys
import time
import urllib.error
import urllib.request
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Generate an image with gpt-image-2 through an OpenAI-compatible /v1/images/generations endpoint."
)
parser.add_argument("--base-url", default=os.environ.get("BASE_URL", "https://claude.omniclaw.store/v1"))
parser.add_argument("--api-key-env", default="API_KEY")
parser.add_argument("--prompt", required=True)
parser.add_argument("--output", default="image.png")
parser.add_argument("--model", default="gpt-image-2")
parser.add_argument("--size", default="1024x1024")
parser.add_argument("--quality", default="medium", choices=["low", "medium", "high", "auto"])
parser.add_argument("--format", default="png", choices=["png", "jpeg", "webp"])
parser.add_argument("--compression", type=int, default=None)
parser.add_argument("--background", default=None, choices=["auto", "opaque"])
parser.add_argument("--moderation", default=None, choices=["auto", "low"])
parser.add_argument("--n", type=int, default=1)
parser.add_argument("--timeout", type=float, default=300.0)
parser.add_argument("--dry-run", action="store_true")
return parser.parse_args()
def build_payload(args: argparse.Namespace) -> dict:
payload = {
"model": args.model,
"prompt": args.prompt,
"size": args.size,
"quality": args.quality,
"output_format": args.format,
"n": args.n,
}
if args.compression is not None:
payload["output_compression"] = args.compression
if args.background:
payload["background"] = args.background
if args.moderation:
payload["moderation"] = args.moderation
return payload
def main() -> int:
args = parse_args()
payload = build_payload(args)
url = args.base_url.rstrip("/") + "/images/generations"
if args.dry_run:
print(json.dumps({"url": url, "payload": payload}, ensure_ascii=False, indent=2))
return 0
api_key = os.environ.get(args.api_key_env)
if not api_key:
print(f"Missing API key env: {args.api_key_env}", file=sys.stderr)
return 2
body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
req = urllib.request.Request(
url,
data=body,
method="POST",
headers={
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
},
)
started = time.monotonic()
try:
with urllib.request.urlopen(req, timeout=args.timeout) as resp:
status = resp.status
raw = resp.read()
except urllib.error.HTTPError as exc:
raw = exc.read()
print(f"HTTP {exc.code}", file=sys.stderr)
print(raw.decode("utf-8", "replace")[:4000], file=sys.stderr)
return 1
except Exception as exc:
print(f"request failed: {type(exc).__name__}: {exc}", file=sys.stderr)
return 1
elapsed = time.monotonic() - started
data = json.loads(raw.decode("utf-8"))
if "error" in data:
print(json.dumps(data["error"], ensure_ascii=False), file=sys.stderr)
return 1
items = data.get("data") or []
if not items or "b64_json" not in items[0]:
print("No b64_json image returned", file=sys.stderr)
print(json.dumps(data, ensure_ascii=False)[:4000], file=sys.stderr)
return 1
image = base64.b64decode(items[0]["b64_json"])
with open(args.output, "wb") as f:
f.write(image)
summary = {
"status": status,
"elapsed_seconds": round(elapsed, 3),
"output": args.output,
"bytes": len(image),
"model": data.get("model"),
"size": data.get("size"),
"quality": data.get("quality"),
"output_format": data.get("output_format"),
"usage": data.get("usage"),
}
print(json.dumps(summary, ensure_ascii=False, indent=2))
return 0
if __name__ == "__main__":
raise SystemExit(main())