#!/usr/bin/env python3
"""Generate / edit images via APIYI's gpt-image-2 series (gpt-image-2 official / gpt-image-2-all / gpt-image-2-vip reverse).
All use the OpenAI Images API (/v1/images/generations + /v1/images/edits), switched by --model. Needs: pip install openai"""
import argparse
import base64
import os
import sys
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from openai import OpenAI
# Max images generated concurrently per call (server returns 1 per call; this simulates more client-side)
MAX_COUNT = 5
# Per-model capability gating: whether these params are accepted (never send the unaccepted ones)
MODEL_CAPS = {
"gpt-image-2": {"size": True, "quality": True, "output_format": True, "mask": True}, # official
"gpt-image-2-all": {"size": False, "quality": False, "output_format": False, "mask": False}, # reverse, ChatGPT
"gpt-image-2-vip": {"size": True, "quality": False, "output_format": False, "mask": False}, # reverse, Codex
}
def caps_of(model):
# Unknown models fall back to the official capability set
return MODEL_CAPS.get(model, MODEL_CAPS["gpt-image-2"])
def load_api_key():
"""Prefer the env var; otherwise look for a .env in the script dir and its parent."""
key = os.environ.get("APIYI_API_KEY")
if key:
return key
here = os.path.dirname(os.path.abspath(__file__))
for d in (here, os.path.dirname(here)):
env_path = os.path.join(d, ".env")
if os.path.exists(env_path):
with open(env_path, encoding="utf-8") as f:
for line in f:
line = line.strip()
if line.startswith("APIYI_API_KEY") and "=" in line:
return line.split("=", 1)[1].strip().strip('"').strip("'")
return None
def project_root():
"""Walk up from the script location to the first dir containing .git or .claude; else cwd."""
d = os.path.dirname(os.path.abspath(__file__))
while True:
if os.path.isdir(os.path.join(d, ".git")) or os.path.isdir(os.path.join(d, ".claude")):
return d
parent = os.path.dirname(d)
if parent == d:
return os.getcwd()
d = parent
def resolve_paths(out, count):
"""Decide output paths. Bare filename -> <project_root>/gpt-image-output/; a path with a dir -> as given."""
if os.path.dirname(out):
base_path = os.path.abspath(out)
else:
out_dir = os.path.join(project_root(), "gpt-image-output")
os.makedirs(out_dir, exist_ok=True)
base_path = os.path.join(out_dir, out)
if count == 1:
return [base_path]
base, ext = os.path.splitext(base_path)
return [f"{base}-{i}{ext}" for i in range(1, count + 1)]
def decode_image(item):
"""Get image bytes: b64_json may be pure base64 or a data:image-prefixed data URL (reverse models); or only a url."""
raw = getattr(item, "b64_json", None)
if raw:
if raw.startswith("data:"):
raw = raw.split(",", 1)[1] # strip the data:image/png;base64, prefix
return base64.b64decode(raw)
url = getattr(item, "url", None)
if url:
with urllib.request.urlopen(url, timeout=360) as r:
return r.read()
raise RuntimeError("response has neither b64_json nor url")
def one_image(client, model, args):
"""Make one request, return image bytes; raise on failure (caught by _safe). Never sends n (default 1; concurrency for more)."""
cap = caps_of(model)
if args.image:
# Edit / multi-image fusion: reopen files each call to avoid sharing handles across threads
files = [open(p, "rb") for p in args.image]
try:
kwargs = dict(model=model, image=files if len(files) > 1 else files[0], prompt=args.prompt)
if cap["size"] and args.size:
kwargs["size"] = args.size
if cap["quality"] and args.quality:
kwargs["quality"] = args.quality
if cap["mask"] and args.mask:
kwargs["mask"] = open(args.mask, "rb")
resp = client.images.edit(**kwargs)
finally:
for fh in files:
fh.close()
else:
# Text to image
kwargs = dict(model=model, prompt=args.prompt)
if cap["size"] and args.size:
kwargs["size"] = args.size
if cap["quality"] and args.quality:
kwargs["quality"] = args.quality
if cap["output_format"] and args.format:
kwargs["output_format"] = args.format
resp = client.images.generate(**kwargs)
return decode_image(resp.data[0])
def main():
api_key = load_api_key()
if not api_key:
sys.exit("No API key found: add a line APIYI_API_KEY=sk-xxx to the .env in the skill folder")
default_model = os.environ.get("APIYI_IMAGE_MODEL", "gpt-image-2")
# Synchronous blocking call; image generation is slow, so give a generous 360s timeout
client = OpenAI(api_key=api_key, base_url="https://api.apiyi.com/v1", timeout=360)
parser = argparse.ArgumentParser(description="gpt-image-2 series image generation")
parser.add_argument("prompt", help="Prompt / edit instruction")
parser.add_argument("--model", default=default_model,
help="gpt-image-2 (official) / gpt-image-2-all (fastest) / gpt-image-2-vip (lockable size)")
parser.add_argument("-i", "--image", action="append", default=[],
help="Input image path (repeatable, up to 16; presence = edit/fusion mode)")
parser.add_argument("-o", "--out", default="output.png", help="Output filename")
parser.add_argument("-n", "--count", type=int, default=1,
help=f"How many at once, default 1, max {MAX_COUNT} (client-side concurrency)")
parser.add_argument("--size", default="auto",
help="Size, e.g. 1024x1024 / 2048x1152 / auto (gpt-image-2-all ignores it; put it in the prompt)")
parser.add_argument("--quality", default="high",
help="Quality low / medium / high / auto (only effective on gpt-image-2 official)")
parser.add_argument("--format", default="png", help="Output format png / jpeg / webp (official only)")
parser.add_argument("--mask", help="Mask image (official edit only, PNG with alpha, applies to the first image)")
args = parser.parse_args()
count = args.count
if count < 1:
count = 1
if count > MAX_COUNT:
print(f"Note: max {MAX_COUNT} at once; clamped {args.count} to {MAX_COUNT}.", file=sys.stderr)
count = MAX_COUNT
paths = resolve_paths(args.out, count)
def task(path):
data = one_image(client, args.model, args)
with open(path, "wb") as f:
f.write(data)
return os.path.abspath(path)
failures = 0
with ThreadPoolExecutor(max_workers=count) as pool:
for path, result in zip(paths, pool.map(lambda p: _safe(task, p), paths)):
ok, value = result
if ok:
print(f"Image saved to {value}")
else:
failures += 1
print(f"Image {os.path.basename(path)} failed: {value}", file=sys.stderr)
if failures == count:
sys.exit("All generations failed.")
def _safe(fn, arg):
try:
return True, fn(arg)
except Exception as e: # noqa: BLE001 — one failure should not abort the other concurrent tasks
return False, str(e)
if __name__ == "__main__":
main()
