添加 PinAI 生图与 Key 设置脚本

支持默认 PinAI 地址、一键配置 Key 和生成元数据输出。
2026-08-09 07:53:40 +02:00
parent 2dd17e5248
commit e8c2c59c40
2 changed files with 359 additions and 0 deletions
@@ -0,0 +1,264 @@
+#!/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())
@@ -0,0 +1,95 @@
+param(
+    [switch]$ValidateOnly
+)
+
+Add-Type -AssemblyName System.Windows.Forms
+Add-Type -AssemblyName System.Drawing
+
+if ($ValidateOnly) {
+    exit 0
+}
+
+$form = New-Object System.Windows.Forms.Form
+$form.Text = '设置 PinAI Key'
+$form.Size = New-Object System.Drawing.Size(540, 250)
+$form.StartPosition = 'CenterScreen'
+$form.FormBorderStyle = 'FixedDialog'
+$form.MaximizeBox = $false
+$form.MinimizeBox = $false
+$form.ShowInTaskbar = $true
+
+$title = New-Object System.Windows.Forms.Label
+$title.Text = '粘贴你的 PinAI API Key'
+$title.Font = New-Object System.Drawing.Font('Microsoft YaHei UI', 12, [System.Drawing.FontStyle]::Bold)
+$title.AutoSize = $true
+$title.Location = New-Object System.Drawing.Point(28, 24)
+$form.Controls.Add($title)
+
+$hint = New-Object System.Windows.Forms.Label
+$hint.Text = 'Key 会以圆点隐藏，并只保存到当前 Windows 用户环境变量。'
+$hint.AutoSize = $true
+$hint.Location = New-Object System.Drawing.Point(30, 58)
+$form.Controls.Add($hint)
+
+$keyBox = New-Object System.Windows.Forms.TextBox
+$keyBox.Location = New-Object System.Drawing.Point(30, 90)
+$keyBox.Size = New-Object System.Drawing.Size(465, 30)
+$keyBox.UseSystemPasswordChar = $true
+$keyBox.Font = New-Object System.Drawing.Font('Segoe UI', 11)
+$form.Controls.Add($keyBox)
+
+$saveButton = New-Object System.Windows.Forms.Button
+$saveButton.Text = '保存 Key'
+$saveButton.Size = New-Object System.Drawing.Size(105, 36)
+$saveButton.Location = New-Object System.Drawing.Point(275, 145)
+$saveButton.DialogResult = [System.Windows.Forms.DialogResult]::None
+$form.Controls.Add($saveButton)
+
+$cancelButton = New-Object System.Windows.Forms.Button
+$cancelButton.Text = '取消'
+$cancelButton.Size = New-Object System.Drawing.Size(105, 36)
+$cancelButton.Location = New-Object System.Drawing.Point(390, 145)
+$cancelButton.DialogResult = [System.Windows.Forms.DialogResult]::Cancel
+$form.Controls.Add($cancelButton)
+
+$form.AcceptButton = $saveButton
+$form.CancelButton = $cancelButton
+$form.Add_Shown({ $keyBox.Focus() })
+
+$saveButton.Add_Click({
+    $key = $keyBox.Text.Trim()
+    if ([string]::IsNullOrWhiteSpace($key)) {
+        [System.Windows.Forms.MessageBox]::Show(
+            '请先粘贴 PinAI API Key。',
+            '未填写 Key',
+            [System.Windows.Forms.MessageBoxButtons]::OK,
+            [System.Windows.Forms.MessageBoxIcon]::Warning
+        ) | Out-Null
+        return
+    }
+
+    try {
+        [Environment]::SetEnvironmentVariable('PINAI_API_KEY', $key, 'User')
+        [Environment]::SetEnvironmentVariable('PINAI_API_KEY', $key, 'Process')
+    }
+    catch {
+        [System.Windows.Forms.MessageBox]::Show(
+            '保存失败。请确认当前 Windows 用户有写入环境变量的权限。',
+            '保存失败',
+            [System.Windows.Forms.MessageBoxButtons]::OK,
+            [System.Windows.Forms.MessageBoxIcon]::Error
+        ) | Out-Null
+        return
+    }
+
+    [System.Windows.Forms.MessageBox]::Show(
+        'PinAI Key 已保存。请完全退出并重新打开 Codex，然后即可生图。',
+        '设置完成',
+        [System.Windows.Forms.MessageBoxButtons]::OK,
+        [System.Windows.Forms.MessageBoxIcon]::Information
+    ) | Out-Null
+    $form.Close()
+})
+
+[void]$form.ShowDialog()
+