90 lines
4.5 KiB
Python
90 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Submit and poll asynchronous Sub2API image tasks."""
|
|
from __future__ import annotations
|
|
import argparse, base64, json, os, sys, time
|
|
from pathlib import Path
|
|
from urllib.request import Request, urlopen
|
|
from urllib.error import HTTPError, URLError
|
|
|
|
ORIGIN = "https://sub2apius.shujk.top"
|
|
AUTH = Path.home() / ".codex/auth.json"
|
|
|
|
def die(msg: str) -> None:
|
|
print(f"Error: {msg}", file=sys.stderr); raise SystemExit(1)
|
|
|
|
def key(path: Path) -> str:
|
|
try: data = json.loads(path.read_text())
|
|
except Exception as exc: die(f"cannot read Codex auth: {exc}")
|
|
value = data.get("OPENAI_API_KEY") if isinstance(data, dict) else None
|
|
if not isinstance(value, str) or not value.strip(): die("OPENAI_API_KEY is missing")
|
|
return value.strip()
|
|
|
|
def request(method: str, url: str, token: str, body: bytes | None = None) -> tuple[int, dict, dict]:
|
|
headers = {"Authorization": f"Bearer {token}", "Accept": "application/json", "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 Chrome/131 Safari/537.36"}
|
|
if body is not None: headers["Content-Type"] = "application/json"
|
|
req = Request(url, data=body, headers=headers, method=method)
|
|
try:
|
|
with urlopen(req, timeout=60) as resp:
|
|
return resp.status, dict(resp.headers), json.loads(resp.read())
|
|
except HTTPError as exc:
|
|
raw = exc.read().decode(errors="replace")
|
|
try: payload = json.loads(raw)
|
|
except json.JSONDecodeError: payload = {"error": raw}
|
|
return exc.code, dict(exc.headers), payload
|
|
except URLError as exc: die(f"request failed: {exc}")
|
|
|
|
def main() -> int:
|
|
p = argparse.ArgumentParser(description="Asynchronous Shujk image generation")
|
|
p.add_argument("--auth-file", type=Path, default=AUTH)
|
|
p.add_argument("--model", default="gpt-image-2")
|
|
p.add_argument("--prompt", required=True)
|
|
p.add_argument("--out", type=Path, required=True)
|
|
p.add_argument("--size", default="1024x1024")
|
|
p.add_argument("--quality", default="low")
|
|
p.add_argument("--poll-seconds", type=float, default=3)
|
|
p.add_argument("--timeout-seconds", type=float, default=1800)
|
|
p.add_argument("--dry-run", action="store_true")
|
|
args = p.parse_args()
|
|
payload = {"model": args.model, "prompt": args.prompt, "size": args.size, "quality": args.quality}
|
|
if args.dry_run:
|
|
print(json.dumps({"url": ORIGIN + "/v1/images/generations/async", "payload": payload}, indent=2)); return 0
|
|
token = key(args.auth_file)
|
|
code, headers, task = request("POST", ORIGIN + "/v1/images/generations/async", token, json.dumps(payload).encode())
|
|
if code != 202: die(f"async submit failed ({code}): {task}")
|
|
task_url = task.get("poll_url") or f"{ORIGIN}/v1/images/tasks/{task.get('task_id')}"
|
|
if task_url.startswith("/"): task_url = ORIGIN + task_url
|
|
deadline = time.monotonic() + args.timeout_seconds
|
|
while time.monotonic() < deadline:
|
|
try:
|
|
code, headers, state = request("GET", task_url, token)
|
|
except TimeoutError:
|
|
time.sleep(args.poll_seconds)
|
|
continue
|
|
if code >= 400: die(f"async poll failed ({code}): {state}")
|
|
status = state.get("status")
|
|
if status == "completed":
|
|
url = state.get("image_url")
|
|
if not url: die("completed task has no image_url")
|
|
data = None
|
|
last_error = None
|
|
for attempt in range(6):
|
|
try:
|
|
with urlopen(Request(url, headers={"Accept": "image/*", "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 Chrome/131 Safari/537.36"}), timeout=120) as resp:
|
|
data = resp.read()
|
|
break
|
|
except HTTPError as exc:
|
|
last_error = exc
|
|
if exc.code != 403 or attempt == 5: break
|
|
time.sleep(2 + attempt)
|
|
except Exception as exc:
|
|
last_error = exc
|
|
break
|
|
if data is None: die(f"image download failed: {last_error}")
|
|
args.out.parent.mkdir(parents=True, exist_ok=True); args.out.write_bytes(data)
|
|
print(json.dumps({"task_id": task.get("task_id"), "status": status, "out": str(args.out), "bytes": len(data)})); return 0
|
|
if status == "failed": die(f"image task failed: {state.get('error')}")
|
|
delay = float(headers.get("Retry-After", args.poll_seconds)); time.sleep(max(1.0, delay))
|
|
die("timed out waiting for image task")
|
|
|
|
if __name__ == "__main__": raise SystemExit(main())
|