265 lines
10 KiB
Python
265 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate a PNG through the default PinAI Images API."""
|
|
|
|
import argparse
|
|
import base64
|
|
import json
|
|
import os
|
|
import re
|
|
import struct
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Mapping, Optional
|
|
from urllib.error import HTTPError, URLError
|
|
from urllib.parse import urlsplit
|
|
from urllib.request import Request, urlopen
|
|
|
|
DEFAULT_API_URL = "https://api.pinaic.com/v1"
|
|
DEFAULT_MODEL = "gpt-image-2"
|
|
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
|
|
PNG_COLOR_TYPES = {0: "grayscale", 2: "RGB", 3: "indexed color", 4: "grayscale with alpha", 6: "RGBA"}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ImageGeneration:
|
|
"""Generated image data plus client-side timing details."""
|
|
|
|
image: bytes
|
|
api_request_seconds: float
|
|
image_download_seconds: float
|
|
response_mode: str
|
|
|
|
|
|
def extract_image_bytes(
|
|
payload: Mapping[str, Any],
|
|
download_url: Optional[Callable[[str], bytes]] = None,
|
|
) -> bytes:
|
|
"""Return the first image as bytes from a standard Images API response."""
|
|
data = payload.get("data")
|
|
if not isinstance(data, list) or not data or not isinstance(data[0], Mapping):
|
|
raise ValueError("API response does not contain data[0]")
|
|
|
|
image = data[0]
|
|
encoded = image.get("b64_json")
|
|
if isinstance(encoded, str) and encoded:
|
|
try:
|
|
return base64.b64decode(encoded, validate=True)
|
|
except (ValueError, TypeError) as error:
|
|
raise ValueError("API response contains invalid b64_json") from error
|
|
|
|
image_url = image.get("url")
|
|
if isinstance(image_url, str) and image_url:
|
|
return (download_url or download_image)(image_url)
|
|
|
|
raise ValueError("API response data[0] must contain b64_json or url")
|
|
|
|
|
|
def download_image(image_url: str, timeout: int = 120) -> bytes:
|
|
parsed = urlsplit(image_url)
|
|
if parsed.scheme != "https" or not parsed.netloc:
|
|
raise ValueError("Image URL must use HTTPS")
|
|
request = Request(image_url, headers={"User-Agent": "openai-imagegen-skill/1.0"})
|
|
try:
|
|
with urlopen(request, timeout=timeout) as response:
|
|
return response.read()
|
|
except (HTTPError, URLError, TimeoutError) as error:
|
|
raise RuntimeError("Could not download image returned by provider") from error
|
|
|
|
|
|
def redact_message(message: str, api_key: str) -> str:
|
|
cleaned = message.replace(api_key, "[REDACTED]") if api_key else message
|
|
return re.sub(
|
|
r"(?i)(authorization|api[_-]?key|token)(\s*[:=]\s*)([^\s,;]+)",
|
|
r"\1\2[REDACTED]",
|
|
cleaned,
|
|
)
|
|
|
|
|
|
def images_endpoint(api_url: str) -> str:
|
|
"""Accept either an API base URL or a complete images/generations endpoint."""
|
|
normalized = api_url.rstrip("/")
|
|
if normalized.endswith("/images/generations"):
|
|
return normalized
|
|
return normalized + "/images/generations"
|
|
|
|
|
|
def request_image_with_timing(
|
|
*,
|
|
prompt: str,
|
|
api_key: str,
|
|
api_url: str,
|
|
model: str,
|
|
size: Optional[str],
|
|
quality: Optional[str],
|
|
timeout: int,
|
|
) -> ImageGeneration:
|
|
body: dict[str, str] = {"model": model, "prompt": prompt, "output_format": "png"}
|
|
if size:
|
|
body["size"] = size
|
|
if quality:
|
|
body["quality"] = quality
|
|
endpoint = images_endpoint(api_url)
|
|
request = Request(
|
|
endpoint,
|
|
data=json.dumps(body).encode("utf-8"),
|
|
headers={
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
"Accept": "application/json",
|
|
},
|
|
method="POST",
|
|
)
|
|
api_started = time.perf_counter()
|
|
try:
|
|
with urlopen(request, timeout=timeout) as response:
|
|
payload = json.loads(response.read().decode("utf-8"))
|
|
except HTTPError as error:
|
|
details = error.read().decode("utf-8", errors="replace")[:1000]
|
|
raise RuntimeError(
|
|
f"Images API returned HTTP {error.code}: {redact_message(details, api_key)}"
|
|
) from error
|
|
except (URLError, TimeoutError) as error:
|
|
raise RuntimeError("Could not reach the Images API") from error
|
|
except (UnicodeDecodeError, json.JSONDecodeError) as error:
|
|
raise RuntimeError("Images API returned invalid JSON") from error
|
|
api_request_seconds = time.perf_counter() - api_started
|
|
|
|
data = payload.get("data")
|
|
first_image = data[0] if isinstance(data, list) and data else {}
|
|
response_mode = "b64_json"
|
|
if isinstance(first_image, Mapping) and not first_image.get("b64_json") and first_image.get("url"):
|
|
response_mode = "url"
|
|
download_started = time.perf_counter()
|
|
image = extract_image_bytes(payload, lambda url: download_image(url, timeout))
|
|
image_download_seconds = time.perf_counter() - download_started if response_mode == "url" else 0.0
|
|
return ImageGeneration(image, api_request_seconds, image_download_seconds, response_mode)
|
|
|
|
|
|
def request_image(
|
|
*, prompt: str, api_key: str, base_url: str, model: str, size: Optional[str], timeout: int
|
|
) -> bytes:
|
|
"""Return generated bytes for callers that do not need timing details."""
|
|
return request_image_with_timing(
|
|
prompt=prompt,
|
|
api_key=api_key,
|
|
api_url=base_url,
|
|
model=model,
|
|
size=size,
|
|
quality=None,
|
|
timeout=timeout,
|
|
).image
|
|
|
|
|
|
def save_png(output: Path, image: bytes, force: bool) -> None:
|
|
if not image.startswith(PNG_SIGNATURE):
|
|
raise ValueError("Provider response is not PNG data")
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
if output.is_symlink():
|
|
raise ValueError("Refusing to write through a symbolic-link output path")
|
|
if output.exists() and not force:
|
|
raise FileExistsError(f"Output already exists: {output}. Use --force to replace it.")
|
|
with tempfile.NamedTemporaryFile(dir=output.parent, delete=False) as handle:
|
|
temporary_path = Path(handle.name)
|
|
handle.write(image)
|
|
try:
|
|
os.replace(temporary_path, output)
|
|
finally:
|
|
temporary_path.unlink(missing_ok=True)
|
|
|
|
|
|
def png_metadata(image: bytes) -> tuple[int, int, int, str]:
|
|
"""Return width, height, bit depth, and color mode from a PNG IHDR chunk."""
|
|
if not image.startswith(PNG_SIGNATURE) or len(image) < 29 or image[12:16] != b"IHDR":
|
|
raise ValueError("Provider response is not a valid PNG with an IHDR chunk")
|
|
width, height, bit_depth, color_type = struct.unpack(">IIBB", image[16:26])
|
|
color_mode = PNG_COLOR_TYPES.get(color_type, f"unknown color type {color_type}")
|
|
return width, height, bit_depth, color_mode
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Generate a PNG through the default PinAI Images API.")
|
|
parser.add_argument("--prompt", required=True, help="Image description")
|
|
parser.add_argument("--output", required=True, type=Path, help="Destination .png file")
|
|
parser.add_argument("--size", help="Optional image size, for example 1024x1024")
|
|
parser.add_argument(
|
|
"--quality",
|
|
default=os.getenv("IMAGE_API_QUALITY") or os.getenv("PINAI_IMAGE_QUALITY"),
|
|
help="Optional provider quality setting, for example low, medium, high, or auto",
|
|
)
|
|
parser.add_argument(
|
|
"--api-url",
|
|
"--base-url",
|
|
dest="api_url",
|
|
default=os.getenv("IMAGE_API_URL") or os.getenv("PINAI_BASE_URL") or DEFAULT_API_URL,
|
|
help="Optional override for the default PinAI API base URL or complete endpoint",
|
|
)
|
|
parser.add_argument(
|
|
"--model",
|
|
default=os.getenv("IMAGE_API_MODEL") or os.getenv("PINAI_IMAGE_MODEL") or DEFAULT_MODEL,
|
|
help=f"Image model (default: {DEFAULT_MODEL})",
|
|
)
|
|
parser.add_argument("--timeout", type=int, default=120, help="Network timeout in seconds")
|
|
parser.add_argument("--force", action="store_true", help="Replace an existing output file")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
api_key = os.getenv("IMAGE_API_KEY") or os.getenv("PINAI_API_KEY")
|
|
if not api_key:
|
|
print("PINAI_API_KEY is required in the environment.", file=sys.stderr)
|
|
return 2
|
|
if not args.api_url:
|
|
print("PinAI API URL is not configured.", file=sys.stderr)
|
|
return 2
|
|
if args.timeout <= 0:
|
|
print("--timeout must be a positive integer.", file=sys.stderr)
|
|
return 2
|
|
if args.output.suffix.lower() != ".png":
|
|
print("--output must use a .png extension.", file=sys.stderr)
|
|
return 2
|
|
started_at = datetime.now(timezone.utc)
|
|
total_started = time.perf_counter()
|
|
try:
|
|
generation = request_image_with_timing(
|
|
prompt=args.prompt,
|
|
api_key=api_key,
|
|
api_url=args.api_url,
|
|
model=args.model,
|
|
size=args.size,
|
|
quality=args.quality,
|
|
timeout=args.timeout,
|
|
)
|
|
width, height, bit_depth, color_mode = png_metadata(generation.image)
|
|
save_started = time.perf_counter()
|
|
save_png(args.output, generation.image, args.force)
|
|
save_seconds = time.perf_counter() - save_started
|
|
except (FileExistsError, RuntimeError, ValueError) as error:
|
|
print(redact_message(str(error), api_key), file=sys.stderr)
|
|
return 1
|
|
total_seconds = time.perf_counter() - total_started
|
|
finished_at = datetime.now(timezone.utc)
|
|
print(f"Saved PNG: {args.output}")
|
|
print(f"Started at (UTC): {started_at.isoformat(timespec='seconds')}")
|
|
print(f"Finished at (UTC): {finished_at.isoformat(timespec='seconds')}")
|
|
print(f"Provider request: {generation.api_request_seconds:.2f}s")
|
|
print(f"Image download: {generation.image_download_seconds:.2f}s")
|
|
print(f"File save: {save_seconds:.2f}s")
|
|
print(f"Total elapsed: {total_seconds:.2f}s")
|
|
print(f"Response mode: {generation.response_mode}")
|
|
print(f"Actual dimensions: {width}x{height}px")
|
|
print(f"Requested size: {args.size or 'provider default'}")
|
|
print(f"Requested quality: {args.quality or 'provider default'}")
|
|
print(f"Image quality: lossless PNG, {bit_depth}-bit {color_mode}")
|
|
file_size = args.output.stat().st_size
|
|
print(f"File size: {file_size} bytes ({file_size / (1024 * 1024):.2f} MiB)")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|