867 lines
34 KiB
Python
Executable File
867 lines
34 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""le-gpt-image 生图 CLI：复用本机 Codex 的 ChatGPT 登录调 GPT Image，无需 OpenAI API Key。
|
||
|
||
用法：
|
||
gpt_image.py generate --prompt "..." --out cover.png --size portrait --quality high
|
||
gpt_image.py edit --image in.png --prompt "把背景换成雪山" --out out.png
|
||
gpt_image.py check # 查看凭证状态（不生图、不打印任何密钥）
|
||
|
||
协议参考 jdmnk/codex-imagegen-cli（Apache-2.0）与 OpenAI Codex CLI 的请求格式：
|
||
读取 ~/.codex/auth.json → access_token 临近过期时用 refresh_token 换新并原子写回 →
|
||
带 Codex 客户端请求头 POST chatgpt.com/backend-api/codex/images/generations|edits。
|
||
仅用 python3 标准库；Pillow 为可选增强（存在时把改图输入压缩成 WebP，否则原样内嵌）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import base64
|
||
import json
|
||
import math
|
||
import os
|
||
import platform
|
||
import re
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
from contextlib import contextmanager
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
from typing import Any, Sequence
|
||
from urllib import error, request
|
||
|
||
try:
|
||
import fcntl
|
||
except ImportError: # Windows：无 flock，退化为无锁刷新
|
||
fcntl = None # type: ignore[assignment]
|
||
|
||
try:
|
||
import tomllib
|
||
except ImportError: # Python < 3.11
|
||
tomllib = None # type: ignore[assignment]
|
||
|
||
__version__ = "1.0.0"
|
||
|
||
DEFAULT_BASE_URL = "https://chatgpt.com/backend-api/codex"
|
||
REFRESH_URL = "https://auth.openai.com/oauth/token"
|
||
CODEX_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
|
||
NATIVE_MODEL = "gpt-image-2" # 请求格式匹配值，实际图像模型由服务端决定
|
||
DEFAULT_REASONING_MODEL = "gpt-5.5"
|
||
ORIGINATOR = "codex_cli_rs"
|
||
MAX_EDIT_IMAGES = 5
|
||
MAX_RETRIES = 4
|
||
TOKEN_LEEWAY_SECONDS = 60
|
||
INPUT_MAX_EDGE = 1536
|
||
|
||
# 尺寸别名 → 具体像素（全部满足 16 倍数 / 边≤3840 / 1:3–3:1 / 总像素上下限）
|
||
SIZE_ALIASES = {
|
||
"1k": "1024x1024",
|
||
"square": "1024x1024",
|
||
"2k": "2048x2048",
|
||
"4k": "3840x2160",
|
||
"portrait": "1024x1536",
|
||
"landscape": "1536x1024",
|
||
"wide": "1536x640",
|
||
"tall": "1024x2048",
|
||
}
|
||
|
||
|
||
class CliError(Exception):
|
||
"""预期内的用户可见失败。"""
|
||
|
||
|
||
class HttpError(CliError):
|
||
def __init__(self, status: int, body: str, retry_after: str | None = None) -> None:
|
||
self.status = status
|
||
self.body = body
|
||
self.retry_after = retry_after
|
||
super().__init__(f"HTTP {status}: {body[:800]}")
|
||
|
||
|
||
def _log(msg: str) -> None:
|
||
print(msg, file=sys.stderr)
|
||
|
||
|
||
def _warn(msg: str) -> None:
|
||
print(f"警告: {msg}", file=sys.stderr)
|
||
|
||
|
||
# ---------------------------------------------------------------- 凭证与令牌
|
||
|
||
def _jwt_payload(token: str) -> dict[str, Any]:
|
||
parts = token.split(".")
|
||
if len(parts) < 2:
|
||
return {}
|
||
try:
|
||
padded = parts[1] + "=" * (-len(parts[1]) % 4)
|
||
payload = json.loads(base64.urlsafe_b64decode(padded.encode("ascii")).decode("utf-8"))
|
||
except (ValueError, UnicodeDecodeError):
|
||
return {}
|
||
return payload if isinstance(payload, dict) else {}
|
||
|
||
|
||
def _token_is_expiring(access_token: str, leeway: int = TOKEN_LEEWAY_SECONDS) -> bool:
|
||
exp = _jwt_payload(access_token).get("exp")
|
||
if not isinstance(exp, (int, float)):
|
||
return False
|
||
return exp <= datetime.now(timezone.utc).timestamp() + leeway
|
||
|
||
|
||
def _codex_home(args: argparse.Namespace) -> Path:
|
||
if getattr(args, "codex_home", None):
|
||
return Path(args.codex_home).expanduser().resolve()
|
||
if os.environ.get("CODEX_HOME"):
|
||
return Path(os.environ["CODEX_HOME"]).expanduser().resolve()
|
||
return Path.home() / ".codex"
|
||
|
||
|
||
def _auth_file(args: argparse.Namespace) -> Path:
|
||
if getattr(args, "auth_file", None):
|
||
return Path(args.auth_file).expanduser().resolve()
|
||
return _codex_home(args) / "auth.json"
|
||
|
||
|
||
def _load_auth(auth_file: Path) -> dict[str, Any]:
|
||
if not auth_file.exists():
|
||
raise CliError(
|
||
f"未找到 Codex 凭证文件: {auth_file}\n"
|
||
"本 skill 走 ChatGPT 订阅通道，需要先在本机登录 Codex CLI（`codex login` 选 ChatGPT 账号），\n"
|
||
'且 Codex 需使用文件凭证存储（config.toml 里 cli_auth_credentials_store = "file"）。'
|
||
)
|
||
try:
|
||
data = json.loads(auth_file.read_text(encoding="utf-8"))
|
||
except (OSError, UnicodeError, json.JSONDecodeError) as exc:
|
||
raise CliError(f"凭证文件不可读或不是合法 JSON: {auth_file} ({exc})") from exc
|
||
if not isinstance(data, dict):
|
||
raise CliError(f"凭证文件内容应为 JSON 对象: {auth_file}")
|
||
return data
|
||
|
||
|
||
def _tokens(auth: dict[str, Any]) -> dict[str, Any]:
|
||
tokens = auth.get("tokens")
|
||
if not isinstance(tokens, dict):
|
||
raise CliError("凭证里没有 ChatGPT tokens，请重新 `codex login`（选 ChatGPT 登录）。")
|
||
return tokens
|
||
|
||
|
||
def _access_token(auth: dict[str, Any]) -> str:
|
||
token = _tokens(auth).get("access_token")
|
||
if not isinstance(token, str) or not token:
|
||
raise CliError("缺少 access_token，请重新 `codex login`。")
|
||
return token
|
||
|
||
|
||
def _id_token_claims(auth: dict[str, Any]) -> dict[str, Any]:
|
||
id_token = _tokens(auth).get("id_token")
|
||
if isinstance(id_token, str) and id_token:
|
||
payload = _jwt_payload(id_token)
|
||
claim = payload.get("https://api.openai.com/auth")
|
||
return claim if isinstance(claim, dict) else {}
|
||
return {}
|
||
|
||
|
||
def _account_id(auth: dict[str, Any]) -> str | None:
|
||
account_id = _tokens(auth).get("account_id")
|
||
if isinstance(account_id, str) and account_id:
|
||
return account_id
|
||
account_id = _id_token_claims(auth).get("chatgpt_account_id")
|
||
return account_id if isinstance(account_id, str) and account_id else None
|
||
|
||
|
||
def _auth_headers(auth: dict[str, Any]) -> dict[str, str]:
|
||
headers = {
|
||
"Authorization": f"Bearer {_access_token(auth)}",
|
||
"Accept": "application/json",
|
||
"Content-Type": "application/json",
|
||
"originator": ORIGINATOR,
|
||
"User-Agent": (
|
||
f"{ORIGINATOR}/{_codex_version()} ({platform.system()} {platform.machine()}) "
|
||
f"le-gpt-image/{__version__}"
|
||
),
|
||
"version": _codex_version(),
|
||
}
|
||
account_id = _account_id(auth)
|
||
if account_id:
|
||
headers["ChatGPT-Account-ID"] = account_id
|
||
if _id_token_claims(auth).get("chatgpt_account_is_fedramp"):
|
||
headers["X-OpenAI-Fedramp"] = "true"
|
||
return headers
|
||
|
||
|
||
def _codex_version() -> str:
|
||
codex_bin = shutil.which("codex")
|
||
if not codex_bin:
|
||
return __version__
|
||
try:
|
||
result = subprocess.run(
|
||
[codex_bin, "--version"], check=False, text=True, capture_output=True, timeout=5
|
||
)
|
||
except (OSError, subprocess.SubprocessError):
|
||
return __version__
|
||
match = re.search(r"(\d+\.\d+\.\d+)", f"{result.stdout}\n{result.stderr}")
|
||
return match.group(1) if match else __version__
|
||
|
||
|
||
def _atomic_write(path: Path, data: bytes, *, replace: bool) -> None:
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".tmp", dir=path.parent)
|
||
try:
|
||
with os.fdopen(fd, "wb") as output:
|
||
output.write(data)
|
||
output.flush()
|
||
os.fsync(output.fileno())
|
||
if replace:
|
||
os.replace(temporary, path)
|
||
else:
|
||
# 原子防覆盖：生成中途出现的输出文件绝不能被悄悄覆盖
|
||
try:
|
||
os.link(temporary, path)
|
||
except FileExistsError as exc:
|
||
raise CliError(f"输出文件已存在: {path}（加 --force 才允许覆盖）") from exc
|
||
finally:
|
||
Path(temporary).unlink(missing_ok=True)
|
||
|
||
|
||
@contextmanager
|
||
def _refresh_lock(auth_file: Path, timeout: float):
|
||
"""与并发实例互斥的刷新锁；锁文件独立于 auth.json，才能扛住原子替换。"""
|
||
lock_path = auth_file.with_name(auth_file.name + ".gptimg.lock")
|
||
fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o600)
|
||
with os.fdopen(fd, "r+b") as lock:
|
||
if fcntl is None:
|
||
yield
|
||
return
|
||
deadline = time.monotonic() + min(timeout, 30.0)
|
||
while True:
|
||
try:
|
||
fcntl.flock(lock.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||
break
|
||
except (BlockingIOError, OSError) as exc:
|
||
if time.monotonic() >= deadline:
|
||
raise CliError("等待另一进程刷新凭证超时，请重试。") from exc
|
||
time.sleep(0.05)
|
||
try:
|
||
yield
|
||
finally:
|
||
fcntl.flock(lock.fileno(), fcntl.LOCK_UN)
|
||
|
||
|
||
def _post_json(
|
||
url: str, *, headers: dict[str, str], payload: dict[str, Any], timeout: float
|
||
) -> dict[str, Any]:
|
||
body = json.dumps(payload).encode("utf-8")
|
||
req = request.Request(url, data=body, headers=headers, method="POST")
|
||
try:
|
||
with request.urlopen(req, timeout=timeout) as response:
|
||
raw = response.read().decode("utf-8")
|
||
except error.HTTPError as exc:
|
||
raw = exc.read().decode("utf-8", errors="replace")
|
||
raise HttpError(exc.code, raw, exc.headers.get("Retry-After") if exc.headers else None) from exc
|
||
except (error.URLError, OSError) as exc:
|
||
raise CliError(f"网络请求失败: {exc}") from exc
|
||
if not raw:
|
||
return {}
|
||
try:
|
||
parsed = json.loads(raw)
|
||
except json.JSONDecodeError as exc:
|
||
raise CliError("响应不是合法 JSON。") from exc
|
||
return parsed if isinstance(parsed, dict) else {}
|
||
|
||
|
||
def _refresh_auth(auth: dict[str, Any], auth_file: Path, *, timeout: float) -> dict[str, Any]:
|
||
import copy
|
||
|
||
with _refresh_lock(auth_file, timeout):
|
||
current = _load_auth(auth_file) if auth_file.exists() else auth
|
||
# 别的进程刚刷新过就直接用新的
|
||
if (
|
||
current is not auth
|
||
and current.get("tokens") != auth.get("tokens")
|
||
and not _token_is_expiring(_access_token(current))
|
||
):
|
||
return current
|
||
updated = copy.deepcopy(current)
|
||
tokens = _tokens(updated)
|
||
refresh_token = tokens.get("refresh_token")
|
||
if not isinstance(refresh_token, str) or not refresh_token:
|
||
raise CliError("缺少 refresh_token，请重新 `codex login`。")
|
||
try:
|
||
refreshed = _post_json(
|
||
REFRESH_URL,
|
||
headers={"Accept": "application/json", "Content-Type": "application/json"},
|
||
payload={
|
||
"client_id": CODEX_CLIENT_ID,
|
||
"grant_type": "refresh_token",
|
||
"refresh_token": refresh_token,
|
||
},
|
||
timeout=timeout,
|
||
)
|
||
except HttpError as exc:
|
||
# OAuth 错误体可能带凭证内容，绝不回显
|
||
raise CliError(f"刷新令牌失败（HTTP {exc.status}），请重新 `codex login`。") from exc
|
||
access = refreshed.get("access_token")
|
||
if not isinstance(access, str) or not access:
|
||
raise CliError("刷新响应里没有 access_token，凭证文件未改动。")
|
||
for key in ("access_token", "refresh_token", "id_token"):
|
||
value = refreshed.get(key)
|
||
if isinstance(value, str) and value:
|
||
tokens[key] = value
|
||
if not isinstance(tokens.get("id_token"), str) or not tokens["id_token"]:
|
||
raise CliError("刷新后 id_token 无效，请重新 `codex login`。")
|
||
updated["last_refresh"] = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
||
_atomic_write(auth_file, (json.dumps(updated, indent=2) + "\n").encode(), replace=True)
|
||
return updated
|
||
|
||
|
||
def _load_ready_auth(args: argparse.Namespace) -> tuple[dict[str, Any], Path]:
|
||
auth_file = _auth_file(args)
|
||
auth = _load_auth(auth_file)
|
||
if _token_is_expiring(_access_token(auth)):
|
||
auth = _refresh_auth(auth, auth_file, timeout=args.timeout)
|
||
return auth, auth_file
|
||
|
||
|
||
# ---------------------------------------------------------------- 请求体
|
||
|
||
def _mime_and_bytes(path: Path) -> tuple[bytes, str]:
|
||
raw = path.read_bytes()
|
||
if raw.startswith(b"\x89PNG\r\n\x1a\n"):
|
||
mime = "image/png"
|
||
elif raw.startswith(b"\xff\xd8\xff"):
|
||
mime = "image/jpeg"
|
||
elif raw[0:4] == b"RIFF" and raw[8:12] == b"WEBP":
|
||
mime = "image/webp"
|
||
elif raw.startswith((b"GIF87a", b"GIF89a")):
|
||
mime = "image/gif"
|
||
else:
|
||
raise CliError(f"无法识别图片格式（支持 PNG/JPEG/WebP/GIF）: {path}")
|
||
return raw, mime
|
||
|
||
|
||
def _data_url_for_image(path: Path, max_edge: int) -> str:
|
||
if not path.is_file():
|
||
raise CliError(f"输入图片不存在: {path}")
|
||
try: # Pillow 可选增强：压到最长边 max_edge 的 WebP，省流量
|
||
from io import BytesIO
|
||
|
||
from PIL import Image, ImageOps # type: ignore[import-not-found]
|
||
|
||
with Image.open(path) as source:
|
||
source.load()
|
||
image = ImageOps.exif_transpose(source)
|
||
width, height = image.size
|
||
edge = max(width, height)
|
||
if 0 < max_edge < edge:
|
||
scale = max_edge / float(edge)
|
||
image = image.resize(
|
||
(max(1, round(width * scale)), max(1, round(height * scale))),
|
||
Image.Resampling.LANCZOS,
|
||
)
|
||
if image.mode not in {"RGB", "RGBA"}:
|
||
image = image.convert("RGBA" if "A" in image.getbands() else "RGB")
|
||
output = BytesIO()
|
||
image.save(output, "WEBP", quality=90)
|
||
raw, mime = output.getvalue(), "image/webp"
|
||
except ImportError:
|
||
raw, mime = _mime_and_bytes(path)
|
||
except OSError as exc:
|
||
raise CliError(f"读取图片失败 {path}: {exc}") from exc
|
||
encoded = base64.b64encode(raw).decode("ascii")
|
||
return f"data:{mime};base64,{encoded}"
|
||
|
||
|
||
STYLE_TRANSFER_PROMPT = (
|
||
"Create a new version of the first input image. Preserve its main subject, identity, "
|
||
"composition, proportions, and important details. Use the final input image only as a style "
|
||
"reference: apply its visual medium, rendering approach, color palette, lighting, texture, "
|
||
"finish, and mood. Do not copy the style reference's subject or scene content."
|
||
)
|
||
|
||
|
||
def _native_payload(
|
||
*, prompt: str, args: argparse.Namespace, mode: str, images: Sequence[Path] | None
|
||
) -> dict[str, Any]:
|
||
payload: dict[str, Any] = {
|
||
"model": NATIVE_MODEL,
|
||
"prompt": prompt,
|
||
"size": args.size,
|
||
"quality": args.quality,
|
||
"background": args.background,
|
||
}
|
||
if mode == "edit":
|
||
if not images or len(images) > MAX_EDIT_IMAGES:
|
||
raise CliError(f"改图需要 1–{MAX_EDIT_IMAGES} 张输入图。")
|
||
payload["images"] = [
|
||
{"image_url": _data_url_for_image(path, args.input_max_edge)} for path in images
|
||
]
|
||
return payload
|
||
|
||
|
||
def _reasoning_model(args: argparse.Namespace) -> str:
|
||
if args.model:
|
||
return args.model
|
||
if os.environ.get("CODEX_IMAGEGEN_MODEL"):
|
||
return os.environ["CODEX_IMAGEGEN_MODEL"]
|
||
config_file = _codex_home(args) / "config.toml"
|
||
if tomllib is not None and config_file.exists():
|
||
try:
|
||
config = tomllib.loads(config_file.read_text(encoding="utf-8"))
|
||
model = config.get("model")
|
||
if isinstance(model, str) and model.strip():
|
||
return model
|
||
except (OSError, UnicodeError, Exception):
|
||
pass
|
||
return DEFAULT_REASONING_MODEL
|
||
|
||
|
||
def _responses_payload(
|
||
*, prompt: str, args: argparse.Namespace, mode: str, images: Sequence[Path] | None
|
||
) -> dict[str, Any]:
|
||
content: list[dict[str, Any]] = [{"type": "input_text", "text": prompt}]
|
||
if images:
|
||
if len(images) > MAX_EDIT_IMAGES:
|
||
raise CliError(f"改图最多 {MAX_EDIT_IMAGES} 张输入图。")
|
||
for image in images:
|
||
content.append(
|
||
{"type": "input_image", "image_url": _data_url_for_image(image, args.input_max_edge)}
|
||
)
|
||
instructions = (
|
||
"Use the available image generation tool to generate exactly one PNG image for the "
|
||
"user request. Do not use any other tool."
|
||
)
|
||
if mode == "edit":
|
||
instructions += " Treat the provided input images as edit/reference images for the request."
|
||
return {
|
||
"model": _reasoning_model(args),
|
||
"instructions": instructions,
|
||
"input": [{"type": "message", "role": "user", "content": content}],
|
||
"tools": [
|
||
{
|
||
"type": "image_generation",
|
||
"output_format": "png",
|
||
"size": args.size,
|
||
"quality": args.quality,
|
||
"background": args.background,
|
||
}
|
||
],
|
||
"tool_choice": {"type": "image_generation"},
|
||
"parallel_tool_calls": False,
|
||
"store": False,
|
||
"stream": True,
|
||
"prompt_cache_key": "le-gpt-image",
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------- 调用与落盘
|
||
|
||
def _native_result(
|
||
url: str, *, headers: dict[str, str], payload: dict[str, Any], timeout: float
|
||
) -> bytes:
|
||
response = _post_json(url, headers=headers, payload=payload, timeout=timeout)
|
||
data = response.get("data")
|
||
if not isinstance(data, list) or len(data) != 1 or not isinstance(data[0], dict):
|
||
error_obj = response.get("error")
|
||
if isinstance(error_obj, dict):
|
||
raise CliError(f"生图失败: {error_obj.get('message') or error_obj}")
|
||
raise CliError("响应里没有图像数据。")
|
||
encoded = data[0].get("b64_json")
|
||
if not isinstance(encoded, str) or not encoded:
|
||
raise CliError("响应里的图像数据为空。")
|
||
return base64.b64decode(encoded)
|
||
|
||
|
||
def _sse_data_lines(url: str, headers: dict[str, str], payload: dict[str, Any], timeout: float):
|
||
req_headers = {**headers, "Accept": "text/event-stream", "Content-Type": "application/json"}
|
||
req = request.Request(
|
||
url, data=json.dumps(payload).encode("utf-8"), headers=req_headers, method="POST"
|
||
)
|
||
try:
|
||
with request.urlopen(req, timeout=timeout) as response:
|
||
buf = ""
|
||
while True:
|
||
chunk = response.read(4096)
|
||
if not chunk:
|
||
break
|
||
buf += chunk.decode("utf-8", "replace")
|
||
while "\n" in buf:
|
||
line, buf = buf.split("\n", 1)
|
||
stripped = line.strip()
|
||
if stripped.startswith("data:"):
|
||
data = stripped[5:].strip()
|
||
if data and data != "[DONE]":
|
||
yield data
|
||
tail = buf.strip()
|
||
if tail.startswith("data:"):
|
||
data = tail[5:].strip()
|
||
if data and data != "[DONE]":
|
||
yield data
|
||
except error.HTTPError as exc:
|
||
raw = exc.read().decode("utf-8", errors="replace")
|
||
raise HttpError(exc.code, raw, exc.headers.get("Retry-After") if exc.headers else None) from exc
|
||
except (error.URLError, OSError) as exc:
|
||
raise CliError(f"网络请求失败: {exc}") from exc
|
||
|
||
|
||
def _responses_result(
|
||
url: str, *, headers: dict[str, str], payload: dict[str, Any], timeout: float
|
||
) -> bytes:
|
||
last_status: str | None = None
|
||
for data in _sse_data_lines(url, headers, payload, timeout):
|
||
try:
|
||
event = json.loads(data)
|
||
except json.JSONDecodeError as exc:
|
||
raise CliError("响应流里包含非法 JSON。") from exc
|
||
if not isinstance(event, dict):
|
||
continue
|
||
if event.get("type") in {"response.failed", "response.incomplete", "error"}:
|
||
error_obj = event.get("error") or (event.get("response") or {}).get("error")
|
||
message = error_obj.get("message") if isinstance(error_obj, dict) else None
|
||
raise CliError(f"responses 通道生图失败: {message or event}")
|
||
items: list[Any] = [event.get("item")]
|
||
if event.get("type") == "response.completed":
|
||
response_obj = event.get("response")
|
||
if isinstance(response_obj, dict) and isinstance(response_obj.get("output"), list):
|
||
items.extend(response_obj["output"])
|
||
for item in items:
|
||
if isinstance(item, dict) and item.get("type") == "image_generation_call":
|
||
last_status = item.get("status")
|
||
result = item.get("result")
|
||
if isinstance(result, str) and result and last_status == "completed":
|
||
return base64.b64decode(result)
|
||
if last_status:
|
||
raise CliError(f"响应流结束但没有拿到图像（最后状态: {last_status}）。")
|
||
raise CliError("响应流里没有图像生成结果。")
|
||
|
||
|
||
def _retry_decision(exc: HttpError, attempt: int) -> float | None:
|
||
"""只在明确限流时重试；传输类/5xx 故意不重放——图像可能已消耗额度。"""
|
||
if exc.status != 429:
|
||
return None
|
||
delay = min(2.0**attempt, 16.0)
|
||
if exc.retry_after:
|
||
try:
|
||
wait = float(exc.retry_after)
|
||
except ValueError:
|
||
wait = 0.0
|
||
delay = max(delay, wait)
|
||
return delay if math.isfinite(delay) and delay <= 300 else None
|
||
|
||
|
||
def _png_size(image_bytes: bytes) -> tuple[int, int] | None:
|
||
if image_bytes[:8] != b"\x89PNG\r\n\x1a\n" or len(image_bytes) < 24:
|
||
return None
|
||
width = int.from_bytes(image_bytes[16:20], "big")
|
||
height = int.from_bytes(image_bytes[20:24], "big")
|
||
return width, height
|
||
|
||
|
||
def _write_image(image_bytes: bytes, out_path: Path, *, force: bool, requested_size: str) -> None:
|
||
actual = _png_size(image_bytes)
|
||
if actual is None:
|
||
raise CliError("服务端返回的不是有效 PNG 图像。")
|
||
_atomic_write(out_path, image_bytes, replace=force)
|
||
if requested_size != "auto":
|
||
requested = requested_size.lower().replace("x", "×")
|
||
actual_text = f"{actual[0]}x{actual[1]}"
|
||
if requested_size != actual_text:
|
||
_warn(f"请求尺寸 {requested_size}，服务端实际返回 {actual_text}，已按实际尺寸保存。")
|
||
_log(f" 已保存 {actual[0]}x{actual[1]} PNG")
|
||
|
||
|
||
def _call_and_save(
|
||
*, args: argparse.Namespace, url: str, payload: dict[str, Any], native: bool, out_path: Path
|
||
) -> None:
|
||
if args.dry_run:
|
||
preview = json.dumps(
|
||
{
|
||
"method": "POST",
|
||
"url": url,
|
||
"headers": {"Authorization": "Bearer <redacted>"},
|
||
"payload": _redact_payload(payload),
|
||
"outputs": [str(out_path)],
|
||
},
|
||
ensure_ascii=False,
|
||
indent=2,
|
||
)
|
||
print(preview)
|
||
return
|
||
auth, auth_file = _load_ready_auth(args)
|
||
headers = _auth_headers(auth)
|
||
call = _native_result if native else _responses_result
|
||
attempt = 0
|
||
refreshed_401 = False
|
||
t0 = time.monotonic()
|
||
while True:
|
||
try:
|
||
image_bytes = call(url, headers=headers, payload=payload, timeout=args.timeout)
|
||
break
|
||
except HttpError as exc:
|
||
if exc.status == 401 and not refreshed_401:
|
||
auth = _refresh_auth(auth, auth_file, timeout=args.timeout)
|
||
headers = _auth_headers(auth)
|
||
refreshed_401 = True
|
||
continue
|
||
delay = _retry_decision(exc, attempt)
|
||
if delay is None or attempt >= MAX_RETRIES:
|
||
raise
|
||
attempt += 1
|
||
_warn(f"被限流，{delay:.0f}s 后重试（{attempt}/{MAX_RETRIES}）")
|
||
time.sleep(delay)
|
||
_write_image(image_bytes, out_path, force=args.force, requested_size=args.size)
|
||
_log(f" 完成（{time.monotonic() - t0:.1f}s）")
|
||
|
||
|
||
def _redact_payload(payload: dict[str, Any]) -> dict[str, Any]:
|
||
def redact(value: Any) -> Any:
|
||
if isinstance(value, dict):
|
||
return {
|
||
key: (f"{str(item).split(';base64,', 1)[0]};base64,<redacted>"
|
||
if key == "image_url" and isinstance(item, str)
|
||
else redact(item))
|
||
for key, item in value.items()
|
||
}
|
||
if isinstance(value, list):
|
||
return [redact(item) for item in value]
|
||
return value
|
||
|
||
return redact(payload)
|
||
|
||
|
||
# ---------------------------------------------------------------- 参数与入口
|
||
|
||
def _parse_size(value: str) -> str:
|
||
value = value.strip().lower()
|
||
if value in SIZE_ALIASES:
|
||
return SIZE_ALIASES[value]
|
||
if value == "auto":
|
||
return value
|
||
match = re.fullmatch(r"([1-9][0-9]{0,3})x([1-9][0-9]{0,3})", value)
|
||
if not match:
|
||
raise argparse.ArgumentTypeError(
|
||
f"尺寸须为 auto、别名（{'/'.join(SIZE_ALIASES)}）或 WIDTHxHEIGHT，收到: {value}"
|
||
)
|
||
width, height = map(int, match.groups())
|
||
if not (
|
||
width % 16 == height % 16 == 0
|
||
and max(width, height) <= 3840
|
||
and max(width, height) <= 3 * min(width, height)
|
||
and 655360 <= width * height <= 8294400
|
||
):
|
||
raise argparse.ArgumentTypeError(
|
||
"尺寸约束不满足: 宽高为 16 的倍数、最长边 ≤3840、宽高比 1:3–3:1、"
|
||
"总像素 655360–8294400。"
|
||
)
|
||
return value
|
||
|
||
|
||
def _slugify(text: str) -> str:
|
||
value = re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")
|
||
return value[:60] or "image"
|
||
|
||
|
||
def _default_out(prompt: str) -> Path:
|
||
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||
return Path("fig") / f"{stamp}-{_slugify(prompt)}.png"
|
||
|
||
|
||
def _resolve_path(raw: str, cd: Path) -> Path:
|
||
path = Path(raw).expanduser()
|
||
if not path.is_absolute():
|
||
path = cd / path
|
||
return path.resolve()
|
||
|
||
|
||
def _read_prompt(args: argparse.Namespace) -> str | None:
|
||
if args.prompt is not None and args.prompt_file is not None:
|
||
raise CliError("--prompt 与 --prompt-file 只能二选一。")
|
||
if args.prompt_file is not None:
|
||
text = Path(args.prompt_file).read_text(encoding="utf-8").strip()
|
||
elif args.prompt is not None:
|
||
text = args.prompt.strip()
|
||
else:
|
||
return None
|
||
if not text:
|
||
raise CliError("prompt 不能为空。")
|
||
return text
|
||
|
||
|
||
def _validate_common(args: argparse.Namespace) -> Path:
|
||
if not math.isfinite(args.timeout) or args.timeout <= 0:
|
||
raise CliError("--timeout 须为正数。")
|
||
if args.n < 1:
|
||
raise CliError("--n 至少为 1。")
|
||
if args.input_max_edge < 0:
|
||
raise CliError("--input-max-edge 须 ≥0。")
|
||
if args.background == "transparent":
|
||
raise CliError("当前 Codex 图像通道不支持透明背景，请用 auto 或 opaque。")
|
||
if args.backend == "native" and getattr(args, "model", None):
|
||
raise CliError("--model 只在 --backend responses 下有效（它选的是推理模型，不是图像模型）。")
|
||
cd = Path(getattr(args, "cd", None) or os.getcwd()).expanduser().resolve()
|
||
if not cd.is_dir():
|
||
raise CliError(f"--cd 目录不存在: {cd}")
|
||
return cd
|
||
|
||
|
||
def _output_paths(out_path: Path, count: int) -> list[Path]:
|
||
if count <= 1:
|
||
return [out_path]
|
||
return [
|
||
out_path.with_name(f"{out_path.stem}-{idx}{out_path.suffix or '.png'}")
|
||
for idx in range(1, count + 1)
|
||
]
|
||
|
||
|
||
def _run_jobs(
|
||
*, args: argparse.Namespace, mode: str, prompt: str, out_path: Path,
|
||
images: Sequence[Path] | None,
|
||
) -> None:
|
||
paths = _output_paths(out_path, args.n)
|
||
if not args.dry_run:
|
||
for path in paths:
|
||
if path.exists() and not args.force:
|
||
raise CliError(f"输出文件已存在: {path}（加 --force 才允许覆盖）")
|
||
count_text = f"{args.n}× " if args.n > 1 else ""
|
||
_log(f"{mode} {count_text}{out_path}")
|
||
if args.backend == "native":
|
||
url = args.base_url.rstrip("/") + ("/images/edits" if mode == "edit" else "/images/generations")
|
||
payload = _native_payload(prompt=prompt, args=args, mode=mode, images=images)
|
||
native = True
|
||
else:
|
||
url = args.base_url.rstrip("/") + "/responses"
|
||
payload = _responses_payload(prompt=prompt, args=args, mode=mode, images=images)
|
||
native = False
|
||
for path in paths:
|
||
_call_and_save(args=args, url=url, payload=payload, native=native, out_path=path)
|
||
print(str(path), flush=True) # stdout 输出最终路径，供 agent 解析
|
||
|
||
|
||
def _cmd_generate(args: argparse.Namespace) -> int:
|
||
cd = _validate_common(args)
|
||
prompt = _read_prompt(args)
|
||
if prompt is None:
|
||
raise CliError("缺少 prompt，用 --prompt 或 --prompt-file。")
|
||
out_path = _resolve_path(args.out, cd) if args.out else (cd / _default_out(prompt))
|
||
_run_jobs(args=args, mode="generate", prompt=prompt, out_path=out_path, images=None)
|
||
return 0
|
||
|
||
|
||
def _cmd_edit(args: argparse.Namespace) -> int:
|
||
cd = _validate_common(args)
|
||
prompt = _read_prompt(args)
|
||
if args.style_image:
|
||
extra = f"\n\nAdditional instruction: {prompt}" if prompt else ""
|
||
prompt = STYLE_TRANSFER_PROMPT + extra
|
||
elif prompt is None:
|
||
raise CliError("缺少 prompt，用 --prompt、--prompt-file 或 --style-image。")
|
||
out_path = _resolve_path(args.out, cd) if args.out else (cd / _default_out(prompt))
|
||
images = [_resolve_path(raw, cd) for raw in args.image]
|
||
if args.style_image:
|
||
images.append(_resolve_path(args.style_image, cd))
|
||
_run_jobs(args=args, mode="edit", prompt=prompt, out_path=out_path, images=images)
|
||
return 0
|
||
|
||
|
||
def _cmd_check(args: argparse.Namespace) -> int:
|
||
auth_file = _auth_file(args)
|
||
auth = _load_auth(auth_file)
|
||
print(f"凭证文件: {auth_file}")
|
||
print(f"登录方式: {auth.get('auth_mode', '未知')}")
|
||
if auth.get("auth_mode") and auth["auth_mode"] != "chatgpt":
|
||
raise CliError("当前不是 ChatGPT 登录态（可能是 API Key 模式），请 `codex login` 选 ChatGPT。")
|
||
tokens = _tokens(auth)
|
||
account_id = _account_id(auth)
|
||
print(f"账号 ID: {account_id[:8]}…（已截断）" if account_id else "账号 ID: 未找到")
|
||
plan = _id_token_claims(auth).get("chatgpt_plan_type")
|
||
print(f"套餐类型: {plan or '未知'}")
|
||
exp = _jwt_payload(_access_token(auth)).get("exp")
|
||
if isinstance(exp, (int, float)):
|
||
hours = (exp - datetime.now(timezone.utc).timestamp()) / 3600
|
||
print(f"access_token: 剩余 {hours:.1f} 小时")
|
||
else:
|
||
print("access_token: 无法解析有效期")
|
||
print(f"refresh_token: {'有' if tokens.get('refresh_token') else '无'}")
|
||
if args.refresh:
|
||
_refresh_auth(auth, auth_file, timeout=args.timeout)
|
||
print("已刷新令牌并写回凭证文件。")
|
||
else:
|
||
print("提示: 加 --refresh 可强制刷新令牌。")
|
||
return 0
|
||
|
||
|
||
def _add_common_args(parser: argparse.ArgumentParser, *, require_out: bool) -> None:
|
||
parser.add_argument("--prompt")
|
||
parser.add_argument("--prompt-file")
|
||
parser.add_argument("--out", required=require_out, help="输出路径；缺省为 fig/时间戳-slug.png")
|
||
parser.add_argument("--force", action="store_true", help="允许覆盖已存在的输出文件")
|
||
parser.add_argument(
|
||
"--size", type=_parse_size, default="auto",
|
||
help="尺寸: auto / 别名(1k 2k 4k square portrait landscape wide tall) / WIDTHxHEIGHT",
|
||
)
|
||
parser.add_argument("--quality", choices=["auto", "low", "medium", "high"], default="auto")
|
||
parser.add_argument("--background", choices=["auto", "opaque"], default="auto")
|
||
parser.add_argument("--n", type=int, default=1, help="生成张数（逐张请求，文件名加 -1..-n）")
|
||
parser.add_argument(
|
||
"--backend", choices=["native", "responses"], default="native",
|
||
help="native=GPT Image 直连通道（默认）；responses=走推理模型的备用通道",
|
||
)
|
||
parser.add_argument("--model", help="responses 通道的推理模型（默认读 Codex 配置或 gpt-5.5）")
|
||
parser.add_argument("--base-url", default=os.environ.get("CODEX_IMAGEGEN_BASE_URL", DEFAULT_BASE_URL))
|
||
parser.add_argument("--timeout", type=float, default=300.0)
|
||
parser.add_argument("--auth-file", help="自定义 auth.json 路径")
|
||
parser.add_argument("--codex-home", help="自定义 Codex 目录（默认 $CODEX_HOME 或 ~/.codex）")
|
||
parser.add_argument("--cd", help="相对路径的基准目录")
|
||
parser.add_argument("--input-max-edge", type=int, default=INPUT_MAX_EDGE,
|
||
help="改图输入压缩的最长边像素，0 为不压缩")
|
||
parser.add_argument("--dry-run", action="store_true", help="只打印请求（凭证打码），不联网不生图")
|
||
|
||
|
||
def build_parser() -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser(
|
||
prog="gpt_image.py",
|
||
description="用本机 Codex ChatGPT 登录调 GPT Image 生图/改图（无需 API Key）",
|
||
)
|
||
parser.add_argument("--version", "-v", action="version", version=f"%(prog)s {__version__}")
|
||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||
|
||
generate = subparsers.add_parser("generate", help="文生图")
|
||
_add_common_args(generate, require_out=False)
|
||
generate.set_defaults(func=_cmd_generate)
|
||
|
||
edit = subparsers.add_parser("edit", help="参考图改图 / 风格迁移")
|
||
_add_common_args(edit, require_out=False)
|
||
edit.add_argument("--image", action="append", required=True,
|
||
help="输入图片路径，可重复（1–5 张）")
|
||
edit.add_argument("--style-image",
|
||
help="风格参考图：内容取 --image，画风取此图（此时 --prompt 变为可选补充）")
|
||
edit.set_defaults(func=_cmd_edit)
|
||
|
||
check = subparsers.add_parser("check", help="检查/刷新本机凭证状态")
|
||
check.add_argument("--refresh", action="store_true", help="强制刷新令牌")
|
||
check.add_argument("--timeout", type=float, default=60.0)
|
||
check.add_argument("--auth-file")
|
||
check.add_argument("--codex-home")
|
||
check.set_defaults(func=_cmd_check)
|
||
|
||
return parser
|
||
|
||
|
||
def main(argv: Sequence[str] | None = None) -> int:
|
||
args = build_parser().parse_args(argv)
|
||
try:
|
||
return int(args.func(args))
|
||
except CliError as exc:
|
||
print(f"错误: {exc}", file=sys.stderr)
|
||
return 1
|
||
except (OSError, UnicodeError) as exc:
|
||
print(f"错误: {exc}", file=sys.stderr)
|
||
return 1
|
||
except KeyboardInterrupt:
|
||
print("已中断。", file=sys.stderr)
|
||
return 130
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|