259 lines
9.4 KiB
Python
259 lines
9.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate an image through an OpenAI-compatible Images API."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import base64
|
|
import binascii
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from typing import Any, Callable
|
|
from urllib.error import HTTPError, URLError
|
|
from urllib.parse import urljoin
|
|
from urllib.parse import urlsplit
|
|
from urllib.request import Request, urlopen
|
|
|
|
MAX_IMAGE_BYTES = 50 * 1024 * 1024
|
|
DOWNLOAD_TIMEOUT_SECONDS = 120
|
|
API_TIMEOUT_SECONDS = 180
|
|
KEYRING_SERVICE = "codex-image-gen-custom"
|
|
KEYRING_USER = "default"
|
|
|
|
|
|
def config_dir() -> Path:
|
|
documents = os.environ.get("CODEX_DOCUMENTS") or str(Path.home() / "Documents")
|
|
return Path(documents) / "Codex" / ".shared" / "image-gen-custom"
|
|
|
|
|
|
def _windows_saved_key_path() -> Path:
|
|
return config_dir() / "api_key.dpapi"
|
|
|
|
|
|
def _decrypt_windows_saved_key(path: Path) -> str:
|
|
shell = shutil.which("powershell.exe") or shutil.which("pwsh")
|
|
if not shell:
|
|
raise RuntimeError("PowerShell is required to read the saved Windows API key")
|
|
script = r"""
|
|
$ErrorActionPreference = 'Stop'
|
|
Add-Type -AssemblyName System.Security
|
|
$protected = [Convert]::FromBase64String((Get-Content -Raw -LiteralPath $env:CODEX_IMAGEGEN_KEY_PATH).Trim())
|
|
$plain = [System.Security.Cryptography.ProtectedData]::Unprotect(
|
|
$protected,
|
|
$null,
|
|
[System.Security.Cryptography.DataProtectionScope]::LocalMachine
|
|
)
|
|
[Text.Encoding]::UTF8.GetString($plain)
|
|
"""
|
|
env = os.environ.copy()
|
|
env["CODEX_IMAGEGEN_KEY_PATH"] = str(path)
|
|
command = [shell, "-NoProfile", "-NonInteractive", "-ExecutionPolicy", "Bypass", "-Command", script]
|
|
result = subprocess.run(command, check=False, capture_output=True, text=True, env=env)
|
|
if result.returncode != 0:
|
|
raise RuntimeError((result.stderr or "PowerShell could not decrypt the saved API key").strip())
|
|
return result.stdout.strip()
|
|
|
|
|
|
def _secure_keyring():
|
|
try:
|
|
import keyring
|
|
except ImportError as exc:
|
|
raise RuntimeError("OPENAI_API_KEY is not set and the keyring package is unavailable") from exc
|
|
backend = keyring.get_keyring()
|
|
backend_id = f"{backend.__class__.__module__}.{backend.__class__.__name__}".lower()
|
|
allowed = ("windows", "macos", "secretservice", "libsecret", "kwallet")
|
|
if not any(marker in backend_id for marker in allowed):
|
|
raise RuntimeError("The configured keyring is not a recognized protected OS credential store")
|
|
return keyring
|
|
|
|
|
|
def _field(item: Any, name: str) -> Any:
|
|
if isinstance(item, dict):
|
|
return item.get(name)
|
|
return getattr(item, name, None)
|
|
|
|
|
|
def _validate_image(raw: bytes) -> None:
|
|
signatures = (
|
|
raw.startswith(b"\x89PNG\r\n\x1a\n"),
|
|
raw.startswith(b"\xff\xd8\xff"),
|
|
raw.startswith((b"GIF87a", b"GIF89a")),
|
|
len(raw) >= 12 and raw[:4] in (b"RIFF", b"RIFX") and raw[8:12] == b"WEBP",
|
|
)
|
|
if not raw or not any(signatures):
|
|
raise ValueError("Image API returned data without a supported image signature")
|
|
|
|
|
|
def _download_image(url: str, *, opener: Callable[..., Any] = urlopen) -> bytes:
|
|
parsed = urlsplit(url)
|
|
if parsed.scheme not in ("http", "https") or not parsed.netloc:
|
|
raise ValueError("Image API returned an invalid image URL")
|
|
request = Request(url, headers={"User-Agent": "codex-image-gen-custom/1"})
|
|
with opener(request, timeout=DOWNLOAD_TIMEOUT_SECONDS) as response:
|
|
content_length = response.headers.get("Content-Length")
|
|
if content_length and int(content_length) > MAX_IMAGE_BYTES:
|
|
raise ValueError(f"Image URL exceeds the {MAX_IMAGE_BYTES}-byte limit")
|
|
raw = response.read(MAX_IMAGE_BYTES + 1)
|
|
if len(raw) > MAX_IMAGE_BYTES:
|
|
raise ValueError(f"Downloaded image exceeds the {MAX_IMAGE_BYTES}-byte limit")
|
|
return raw
|
|
|
|
|
|
def extract_image_bytes(item: Any, *, opener: Callable[..., Any] = urlopen) -> bytes:
|
|
image_b64 = _field(item, "b64_json")
|
|
if image_b64:
|
|
try:
|
|
raw = base64.b64decode(image_b64, validate=True)
|
|
except (binascii.Error, TypeError, ValueError) as exc:
|
|
raise ValueError("Image API returned invalid b64_json") from exc
|
|
else:
|
|
image_url = _field(item, "url")
|
|
if not image_url:
|
|
raise ValueError("Image API response contained neither b64_json nor url")
|
|
raw = _download_image(str(image_url), opener=opener)
|
|
_validate_image(raw)
|
|
return raw
|
|
|
|
|
|
def write_image(raw: bytes, output: Path, *, force: bool) -> None:
|
|
if output.exists() and not force:
|
|
raise FileExistsError(f"Output already exists: {output} (use --force to overwrite)")
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
fd, temporary_name = tempfile.mkstemp(prefix=f".{output.name}.", suffix=".tmp", dir=output.parent)
|
|
temporary = Path(temporary_name)
|
|
try:
|
|
with os.fdopen(fd, "wb") as stream:
|
|
stream.write(raw)
|
|
stream.flush()
|
|
os.fsync(stream.fileno())
|
|
temporary.replace(output)
|
|
finally:
|
|
temporary.unlink(missing_ok=True)
|
|
|
|
|
|
def load_api_key() -> str:
|
|
environment_key = os.environ.get("OPENAI_API_KEY", "").strip()
|
|
if environment_key:
|
|
return environment_key
|
|
if os.name == "nt":
|
|
key_path = _windows_saved_key_path()
|
|
if not key_path.exists():
|
|
raise RuntimeError("No API key found. Run configure or set OPENAI_API_KEY")
|
|
key = _decrypt_windows_saved_key(key_path)
|
|
if not key:
|
|
raise RuntimeError("No API key found. Run configure or set OPENAI_API_KEY")
|
|
return key
|
|
try:
|
|
keyring = _secure_keyring()
|
|
key = keyring.get_password(KEYRING_SERVICE, KEYRING_USER)
|
|
except Exception as exc:
|
|
raise RuntimeError("The OS keyring is unavailable; set OPENAI_API_KEY for this run") from exc
|
|
if not key:
|
|
raise RuntimeError("No API key found. Run configure or set OPENAI_API_KEY")
|
|
return key
|
|
|
|
|
|
def _configured_base_url() -> str:
|
|
environment_url = os.environ.get("OPENAI_BASE_URL", "").strip()
|
|
if environment_url:
|
|
return environment_url
|
|
config_path = config_dir() / "base_url.txt"
|
|
try:
|
|
configured_url = config_path.read_text(encoding="utf-8-sig").strip()
|
|
except OSError:
|
|
configured_url = ""
|
|
return configured_url
|
|
|
|
|
|
def _base_url(override: str = "") -> str:
|
|
base_url = override.strip() or _configured_base_url() or "https://api.openai.com/v1"
|
|
base_url = base_url.rstrip("/")
|
|
if not base_url.lower().endswith("/v1"):
|
|
base_url += "/v1"
|
|
return base_url
|
|
|
|
|
|
def _request_generation(*, api_key: str, base_url: str, payload: dict[str, Any]) -> dict[str, Any]:
|
|
endpoint = urljoin(f"{base_url.rstrip('/')}/", "images/generations")
|
|
request = Request(
|
|
endpoint,
|
|
data=json.dumps(payload).encode("utf-8"),
|
|
headers={
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
"Accept": "application/json",
|
|
"User-Agent": "codex-image-gen-custom/1",
|
|
},
|
|
method="POST",
|
|
)
|
|
try:
|
|
with urlopen(request, timeout=API_TIMEOUT_SECONDS) as response:
|
|
raw = response.read()
|
|
except HTTPError as exc:
|
|
detail = exc.read(4096).decode("utf-8", errors="replace").strip()
|
|
suffix = f": {detail}" if detail else ""
|
|
raise RuntimeError(f"Image API returned HTTP {exc.code}{suffix}") from exc
|
|
except URLError as exc:
|
|
raise RuntimeError(f"Could not reach the Image API: {exc.reason}") from exc
|
|
try:
|
|
result = json.loads(raw.decode("utf-8"))
|
|
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
|
raise ValueError("Image API returned invalid JSON") from exc
|
|
if not isinstance(result, dict):
|
|
raise ValueError("Image API returned a non-object JSON response")
|
|
return result
|
|
|
|
|
|
def generate(args: argparse.Namespace) -> None:
|
|
api_key = load_api_key()
|
|
base_url = _base_url(args.base_url)
|
|
started = time.time()
|
|
print("Calling Image API. This can take up to a couple of minutes.", file=sys.stderr)
|
|
result = _request_generation(
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
payload={
|
|
"model": args.model,
|
|
"prompt": args.prompt,
|
|
"size": args.size,
|
|
"quality": args.quality,
|
|
},
|
|
)
|
|
print(f"Generation completed in {time.time() - started:.1f}s.", file=sys.stderr)
|
|
data = result.get("data")
|
|
if not isinstance(data, list) or not data:
|
|
raise ValueError("Image API returned an empty data list")
|
|
write_image(extract_image_bytes(data[0]), Path(args.out), force=args.force)
|
|
print(f"Wrote {args.out}")
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--base-url", default="", help="OpenAI-compatible API root")
|
|
parser.add_argument("--model", default="gpt-image-2")
|
|
parser.add_argument("--prompt", required=True)
|
|
parser.add_argument("--size", default="1024x1024")
|
|
parser.add_argument("--quality", choices=("low", "medium", "high", "auto"), default="low")
|
|
parser.add_argument("--out", required=True)
|
|
parser.add_argument("--force", action="store_true")
|
|
return parser
|
|
|
|
|
|
def main() -> int:
|
|
try:
|
|
generate(build_parser().parse_args())
|
|
except Exception as exc:
|
|
print(f"Error: {exc}", file=sys.stderr)
|
|
return 1
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|