图片结果与保存

import base64
import json
import sys
import urllib.parse
import urllib.request
from pathlib import Path

if len(sys.argv) != 2:
    raise SystemExit("用法：python3 save_images.py generation.json")

payload = json.loads(Path(sys.argv[1]).read_text(encoding="utf-8"))
if payload.get("error"):
    raise SystemExit("响应包含错误，请先排查 API 调用")
items = payload.get("data")
if not isinstance(items, list) or not items:
    raise SystemExit("响应中没有图片结果")

for index, item in enumerate(items, start=1):
    if item.get("b64_json"):
        encoded = item["b64_json"]
        if encoded.startswith("data:"):
            encoded = encoded.split(",", 1)[1]
        encoded = "".join(encoded.split())
        encoded += "=" * (-len(encoded) % 4)
        raw = base64.b64decode(encoded, validate=True)
    elif item.get("url"):
        url = item["url"]
        if urllib.parse.urlsplit(url).scheme not in ("https", "http"):
            raise SystemExit("拒绝非 HTTP(S) 图片地址")
        with urllib.request.urlopen(url, timeout=60) as download:
            raw = download.read()
    else:
        raise SystemExit(f"第 {index} 项没有图片内容或地址")

    if raw.startswith(b"\x89PNG\r\n\x1a\n"):
        extension = "png"
    elif raw.startswith(b"\xff\xd8\xff"):
        extension = "jpg"
    elif raw[:4] == b"RIFF" and raw[8:12] == b"WEBP":
        extension = "webp"
    else:
        raise SystemExit(f"第 {index} 项不是可识别的 PNG/JPEG/WebP 图片")

    output = Path(f"image-{index}.{extension}")
    if output.exists():
        raise SystemExit(f"{output} 已存在，请更换工作目录或先备份文件")
    output.write_bytes(raw)
    print(f"已保存 {output}（{len(raw)} bytes）")