613 lines
19 KiB
Python
613 lines
19 KiB
Python
#!/usr/bin/env python3
|
||
"""gpt-image-2.5 文生图 / 图生图工具。
|
||
|
||
支持：
|
||
- AI / Agent CLI 调用（参数或 JSON）
|
||
- 流程回调（HTTP webhook / 状态文件 / NDJSON 事件流）
|
||
- 指定本地生成保存路径
|
||
|
||
接口文档: Doc/gpt-image_doc.md
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import base64
|
||
import json
|
||
import mimetypes
|
||
import os
|
||
import sys
|
||
import time
|
||
import urllib.error
|
||
import urllib.parse
|
||
import urllib.request
|
||
from pathlib import Path
|
||
from typing import Any, Callable, Dict, List, Optional, Sequence, Union
|
||
|
||
DEFAULT_MODEL = "gpt-image-2.5"
|
||
VIP_MODEL = "gpt-image-2-vip"
|
||
DEFAULT_API_KEY = "sk-b69877bd82c6412ca4872351f428d580"
|
||
DEFAULT_BASE_URL = "https://grsai.dakka.com.cn" # 国内节点；全球: https://grsaiapi.com
|
||
DEFAULT_ASPECT_RATIO = "1024x1024"
|
||
DEFAULT_QUALITY = "auto"
|
||
DEFAULT_REPLY_TYPE = "json"
|
||
POLL_INTERVAL_SEC = 2.0
|
||
POLL_TIMEOUT_SEC = 300.0
|
||
# 标准模型仅支持约 1K；任一边超过该阈值需改用 VIP 像素尺寸
|
||
MAX_STANDARD_EDGE = 1024
|
||
VIP_MODELS = frozenset(
|
||
{
|
||
"gpt-image-2-vip",
|
||
"gpt-image-2.5-flare",
|
||
"gpt-image-2.5-sunburst",
|
||
}
|
||
)
|
||
|
||
|
||
class GenerateError(RuntimeError):
|
||
"""生成失败。"""
|
||
|
||
|
||
def _now_iso() -> str:
|
||
return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
||
|
||
|
||
def _resolve_api_key(explicit: Optional[str] = None) -> str:
|
||
return (
|
||
explicit
|
||
or os.environ.get("GRSAI_API_KEY")
|
||
or os.environ.get("GPT_IMAGE_API_KEY")
|
||
or DEFAULT_API_KEY
|
||
)
|
||
|
||
|
||
def _normalize_base_url(base_url: str) -> str:
|
||
return base_url.rstrip("/")
|
||
|
||
|
||
def _guess_mime(path: Path) -> str:
|
||
mime, _ = mimetypes.guess_type(str(path))
|
||
return mime or "application/octet-stream"
|
||
|
||
|
||
def _file_to_data_url(path: Path) -> str:
|
||
data = path.read_bytes()
|
||
b64 = base64.b64encode(data).decode("ascii")
|
||
return f"data:{_guess_mime(path)};base64,{b64}"
|
||
|
||
|
||
def parse_pixel_size(aspect_ratio: str) -> Optional[tuple[int, int]]:
|
||
"""解析 '2048x2048' / '2048X1152'；比例字符串（如 16:9）返回 None。"""
|
||
text = (aspect_ratio or "").strip().lower().replace(" ", "")
|
||
if "x" not in text:
|
||
return None
|
||
left, _, right = text.partition("x")
|
||
if not left.isdigit() or not right.isdigit():
|
||
return None
|
||
return int(left), int(right)
|
||
|
||
|
||
def needs_vip_for_size(aspect_ratio: str) -> bool:
|
||
"""任一边 >1024 的像素尺寸需要 VIP 模型。"""
|
||
size = parse_pixel_size(aspect_ratio)
|
||
if size is None:
|
||
return False
|
||
w, h = size
|
||
return w > MAX_STANDARD_EDGE or h > MAX_STANDARD_EDGE
|
||
|
||
|
||
def resolve_model_for_size(model: str, aspect_ratio: str) -> str:
|
||
"""大于 1K 时：已是 VIP 系列则保留；否则自动升级到 gpt-image-2-vip。"""
|
||
if not needs_vip_for_size(aspect_ratio):
|
||
return model
|
||
if model in VIP_MODELS:
|
||
return model
|
||
return VIP_MODEL
|
||
|
||
|
||
def normalize_image_ref(ref: str) -> str:
|
||
"""将本地路径转为 data URL；URL / base64 / data URL 原样返回。"""
|
||
text = ref.strip()
|
||
if not text:
|
||
raise GenerateError("空的参考图路径")
|
||
if text.startswith(("http://", "https://", "data:")):
|
||
return text
|
||
# 纯 base64（无 data: 前缀）
|
||
if len(text) > 256 and all(
|
||
c.isalnum() or c in "+/=\n\r" for c in text[:200]
|
||
) and not Path(text).exists():
|
||
return text
|
||
|
||
path = Path(text).expanduser().resolve()
|
||
if not path.is_file():
|
||
raise GenerateError(f"参考图不存在: {path}")
|
||
return _file_to_data_url(path)
|
||
|
||
|
||
def http_json(
|
||
method: str,
|
||
url: str,
|
||
*,
|
||
headers: Optional[Dict[str, str]] = None,
|
||
body: Optional[Dict[str, Any]] = None,
|
||
timeout: float = 300.0,
|
||
) -> Dict[str, Any]:
|
||
data = None
|
||
req_headers = {"Accept": "application/json", **(headers or {})}
|
||
if body is not None:
|
||
data = json.dumps(body, ensure_ascii=False).encode("utf-8")
|
||
req_headers["Content-Type"] = "application/json"
|
||
|
||
request = urllib.request.Request(url, data=data, headers=req_headers, method=method)
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=timeout) as resp:
|
||
raw = resp.read().decode("utf-8")
|
||
if not raw.strip():
|
||
return {}
|
||
return json.loads(raw)
|
||
except urllib.error.HTTPError as exc:
|
||
err_body = exc.read().decode("utf-8", errors="replace")
|
||
try:
|
||
payload = json.loads(err_body) if err_body else {}
|
||
except json.JSONDecodeError:
|
||
payload = {"error": err_body}
|
||
message = payload.get("error") or err_body or str(exc)
|
||
raise GenerateError(f"HTTP {exc.code}: {message}") from exc
|
||
except urllib.error.URLError as exc:
|
||
raise GenerateError(f"网络错误: {exc}") from exc
|
||
|
||
|
||
def download_file(url: str, dest: Path, timeout: float = 120.0) -> Path:
|
||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||
request = urllib.request.Request(url, headers={"User-Agent": "AIGame-gpt-image/1.0"})
|
||
with urllib.request.urlopen(request, timeout=timeout) as resp:
|
||
dest.write_bytes(resp.read())
|
||
return dest
|
||
|
||
|
||
class CallbackHub:
|
||
"""向 stdout / 文件 / HTTP webhook 广播流程事件。"""
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
callback_url: Optional[str] = None,
|
||
callback_file: Optional[Union[str, Path]] = None,
|
||
quiet: bool = False,
|
||
) -> None:
|
||
self.callback_url = callback_url
|
||
self.callback_file = Path(callback_file) if callback_file else None
|
||
self.quiet = quiet
|
||
self.events: List[Dict[str, Any]] = []
|
||
|
||
def emit(self, event: str, **payload: Any) -> Dict[str, Any]:
|
||
record = {
|
||
"event": event,
|
||
"ts": _now_iso(),
|
||
**payload,
|
||
}
|
||
self.events.append(record)
|
||
|
||
if not self.quiet:
|
||
# NDJSON，便于 AI / 流程引擎按行解析
|
||
print(json.dumps(record, ensure_ascii=False), flush=True)
|
||
|
||
if self.callback_file:
|
||
self.callback_file.parent.mkdir(parents=True, exist_ok=True)
|
||
self.callback_file.write_text(
|
||
json.dumps(
|
||
{"latest": record, "history": self.events},
|
||
ensure_ascii=False,
|
||
indent=2,
|
||
),
|
||
encoding="utf-8",
|
||
)
|
||
|
||
if self.callback_url:
|
||
try:
|
||
http_json("POST", self.callback_url, body=record, timeout=30.0)
|
||
except Exception as exc: # noqa: BLE001 — 回调失败不阻断主流程
|
||
if not self.quiet:
|
||
warn = {
|
||
"event": "callback_error",
|
||
"ts": _now_iso(),
|
||
"error": str(exc),
|
||
"callback_url": self.callback_url,
|
||
}
|
||
print(json.dumps(warn, ensure_ascii=False), flush=True)
|
||
|
||
return record
|
||
|
||
|
||
def query_result(
|
||
base_url: str,
|
||
task_id: str,
|
||
api_key: str,
|
||
timeout: float = 60.0,
|
||
) -> Dict[str, Any]:
|
||
url = f"{_normalize_base_url(base_url)}/v1/api/result?{urllib.parse.urlencode({'id': task_id})}"
|
||
return http_json(
|
||
"GET",
|
||
url,
|
||
headers={"Authorization": f"Bearer {api_key}"},
|
||
timeout=timeout,
|
||
)
|
||
|
||
|
||
def submit_generate(
|
||
*,
|
||
prompt: str,
|
||
api_key: str,
|
||
base_url: str = DEFAULT_BASE_URL,
|
||
model: str = DEFAULT_MODEL,
|
||
images: Optional[Sequence[str]] = None,
|
||
aspect_ratio: str = DEFAULT_ASPECT_RATIO,
|
||
quality: str = DEFAULT_QUALITY,
|
||
reply_type: str = DEFAULT_REPLY_TYPE,
|
||
background: Optional[str] = None,
|
||
mask: Optional[str] = None,
|
||
) -> Dict[str, Any]:
|
||
image_refs = [normalize_image_ref(x) for x in (images or [])]
|
||
body: Dict[str, Any] = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"images": image_refs,
|
||
"aspectRatio": aspect_ratio,
|
||
"quality": quality,
|
||
"replyType": reply_type,
|
||
}
|
||
if background:
|
||
body["background"] = background
|
||
if mask:
|
||
body["mask"] = normalize_image_ref(mask)
|
||
|
||
url = f"{_normalize_base_url(base_url)}/v1/api/generate"
|
||
return http_json(
|
||
"POST",
|
||
url,
|
||
headers={"Authorization": f"Bearer {api_key}"},
|
||
body=body,
|
||
timeout=300.0,
|
||
)
|
||
|
||
|
||
def wait_until_done(
|
||
*,
|
||
base_url: str,
|
||
api_key: str,
|
||
initial: Dict[str, Any],
|
||
reply_type: str,
|
||
callbacks: CallbackHub,
|
||
poll_interval: float = POLL_INTERVAL_SEC,
|
||
poll_timeout: float = POLL_TIMEOUT_SEC,
|
||
) -> Dict[str, Any]:
|
||
status = (initial.get("status") or "").lower()
|
||
if status in {"succeeded", "failed", "violation"} or (
|
||
reply_type == "json" and initial.get("results")
|
||
):
|
||
return initial
|
||
|
||
task_id = initial.get("id")
|
||
if not task_id:
|
||
raise GenerateError(f"响应缺少任务 id: {initial}")
|
||
|
||
deadline = time.time() + poll_timeout
|
||
last_progress = initial.get("progress")
|
||
current = initial
|
||
|
||
while time.time() < deadline:
|
||
current = query_result(base_url, task_id, api_key)
|
||
status = (current.get("status") or "").lower()
|
||
progress = current.get("progress")
|
||
if progress != last_progress:
|
||
callbacks.emit(
|
||
"progress",
|
||
id=task_id,
|
||
status=status,
|
||
progress=progress,
|
||
)
|
||
last_progress = progress
|
||
|
||
if status == "succeeded" or current.get("results"):
|
||
return current
|
||
if status in {"failed", "violation"}:
|
||
raise GenerateError(current.get("error") or f"任务失败: {current}")
|
||
|
||
time.sleep(poll_interval)
|
||
|
||
raise GenerateError(f"等待超时({poll_timeout}s)，最后状态: {current}")
|
||
|
||
|
||
def resolve_output_paths(
|
||
output: Optional[str],
|
||
count: int,
|
||
) -> List[Path]:
|
||
if not output:
|
||
stamp = time.strftime("%Y%m%d_%H%M%S")
|
||
output = str(Path("output") / f"gpt_image_{stamp}.png")
|
||
|
||
path = Path(output).expanduser()
|
||
if count <= 1:
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
return [path]
|
||
|
||
# 多图：若给的是文件路径，则追加 _1/_2；若是目录则编号保存
|
||
if path.suffix:
|
||
stem, suffix = path.stem, path.suffix
|
||
parent = path.parent
|
||
parent.mkdir(parents=True, exist_ok=True)
|
||
return [parent / f"{stem}_{i + 1}{suffix}" for i in range(count)]
|
||
|
||
path.mkdir(parents=True, exist_ok=True)
|
||
return [path / f"image_{i + 1}.png" for i in range(count)]
|
||
|
||
|
||
def save_results(
|
||
result: Dict[str, Any],
|
||
output: Optional[str],
|
||
) -> List[Dict[str, Any]]:
|
||
results = result.get("results") or []
|
||
if not results:
|
||
raise GenerateError(f"成功响应但无 results: {result}")
|
||
|
||
paths = resolve_output_paths(output, len(results))
|
||
saved: List[Dict[str, Any]] = []
|
||
for item, dest in zip(results, paths):
|
||
url = item.get("url")
|
||
if not url:
|
||
raise GenerateError(f"结果缺少 url: {item}")
|
||
download_file(url, dest)
|
||
saved.append({"url": url, "path": str(dest.resolve())})
|
||
return saved
|
||
|
||
|
||
def generate(
|
||
*,
|
||
prompt: str,
|
||
output: Optional[str] = None,
|
||
images: Optional[Sequence[str]] = None,
|
||
model: str = DEFAULT_MODEL,
|
||
api_key: Optional[str] = None,
|
||
base_url: str = DEFAULT_BASE_URL,
|
||
aspect_ratio: str = DEFAULT_ASPECT_RATIO,
|
||
quality: str = DEFAULT_QUALITY,
|
||
reply_type: str = DEFAULT_REPLY_TYPE,
|
||
background: Optional[str] = None,
|
||
mask: Optional[str] = None,
|
||
callback_url: Optional[str] = None,
|
||
callback_file: Optional[str] = None,
|
||
poll_interval: float = POLL_INTERVAL_SEC,
|
||
poll_timeout: float = POLL_TIMEOUT_SEC,
|
||
quiet: bool = False,
|
||
on_event: Optional[Callable[[Dict[str, Any]], None]] = None,
|
||
) -> Dict[str, Any]:
|
||
"""执行一次文/图生图，返回结构化结果（可供流程编排直接使用）。"""
|
||
key = _resolve_api_key(api_key)
|
||
model = resolve_model_for_size(model, aspect_ratio)
|
||
callbacks = CallbackHub(
|
||
callback_url=callback_url,
|
||
callback_file=callback_file,
|
||
quiet=quiet,
|
||
)
|
||
if on_event:
|
||
original_emit = callbacks.emit
|
||
|
||
def emit_with_hook(event: str, **payload: Any) -> Dict[str, Any]:
|
||
record = original_emit(event, **payload)
|
||
on_event(record)
|
||
return record
|
||
|
||
callbacks.emit = emit_with_hook # type: ignore[method-assign]
|
||
|
||
callbacks.emit(
|
||
"started",
|
||
model=model,
|
||
prompt=prompt,
|
||
output=output,
|
||
aspect_ratio=aspect_ratio,
|
||
image_count=len(images or []),
|
||
reply_type=reply_type,
|
||
)
|
||
|
||
try:
|
||
submitted = submit_generate(
|
||
prompt=prompt,
|
||
api_key=key,
|
||
base_url=base_url,
|
||
model=model,
|
||
images=images,
|
||
aspect_ratio=aspect_ratio,
|
||
quality=quality,
|
||
reply_type=reply_type,
|
||
background=background,
|
||
mask=mask,
|
||
)
|
||
callbacks.emit(
|
||
"submitted",
|
||
id=submitted.get("id"),
|
||
status=submitted.get("status"),
|
||
progress=submitted.get("progress"),
|
||
)
|
||
|
||
# async / 未完成时轮询
|
||
effective_reply = reply_type
|
||
if (submitted.get("status") or "").lower() == "running" and not submitted.get(
|
||
"results"
|
||
):
|
||
effective_reply = "async"
|
||
|
||
final = wait_until_done(
|
||
base_url=base_url,
|
||
api_key=key,
|
||
initial=submitted,
|
||
reply_type=effective_reply,
|
||
callbacks=callbacks,
|
||
poll_interval=poll_interval,
|
||
poll_timeout=poll_timeout,
|
||
)
|
||
|
||
status = (final.get("status") or "").lower()
|
||
if status in {"failed", "violation"}:
|
||
raise GenerateError(final.get("error") or f"任务失败: {final}")
|
||
|
||
saved = save_results(final, output)
|
||
result = {
|
||
"ok": True,
|
||
"id": final.get("id"),
|
||
"status": final.get("status") or "succeeded",
|
||
"progress": final.get("progress", 100),
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"files": saved,
|
||
"raw": final,
|
||
}
|
||
callbacks.emit("succeeded", **{k: v for k, v in result.items() if k != "raw"})
|
||
return result
|
||
except Exception as exc:
|
||
err = {"ok": False, "error": str(exc)}
|
||
callbacks.emit("failed", **err)
|
||
raise
|
||
|
||
|
||
def _load_json_arg(raw: Optional[str]) -> Dict[str, Any]:
|
||
if not raw:
|
||
return {}
|
||
if raw == "-":
|
||
return json.load(sys.stdin)
|
||
path = Path(raw)
|
||
if path.is_file():
|
||
return json.loads(path.read_text(encoding="utf-8"))
|
||
return json.loads(raw)
|
||
|
||
|
||
def build_parser() -> argparse.ArgumentParser:
|
||
p = argparse.ArgumentParser(
|
||
prog="gpt_image_generate",
|
||
description="gpt-image-2.5 文生图/图生图（支持 AI 调用与流程回调）",
|
||
)
|
||
p.add_argument("--prompt", "-p", help="提示词")
|
||
p.add_argument(
|
||
"--image",
|
||
"-i",
|
||
action="append",
|
||
dest="images",
|
||
default=[],
|
||
help="参考图（可多次；本地路径 / URL / base64）",
|
||
)
|
||
p.add_argument(
|
||
"--output",
|
||
"-o",
|
||
help="生成图保存路径（文件或目录）。多图时自动编号",
|
||
)
|
||
p.add_argument(
|
||
"--model",
|
||
default=DEFAULT_MODEL,
|
||
help=f"模型，默认 {DEFAULT_MODEL}；像素任一边>1024 时自动升为 {VIP_MODEL}",
|
||
)
|
||
p.add_argument("--api-key", default=None, help="API Key，默认读环境变量或内置配置")
|
||
p.add_argument(
|
||
"--base-url",
|
||
default=DEFAULT_BASE_URL,
|
||
help="API 节点，默认国内节点",
|
||
)
|
||
p.add_argument(
|
||
"--aspect-ratio",
|
||
default=DEFAULT_ASPECT_RATIO,
|
||
help='比例或像素，如 "16:9" / "1024x1024"；>1K 请用像素如 "2048x2048"（自动 VIP）',
|
||
)
|
||
p.add_argument("--quality", default=DEFAULT_QUALITY, help="质量: auto/low/medium/high/...")
|
||
p.add_argument(
|
||
"--reply-type",
|
||
default=DEFAULT_REPLY_TYPE,
|
||
choices=["json", "stream", "async"],
|
||
help="回复类型",
|
||
)
|
||
p.add_argument("--background", default=None, help="背景参数，如 transparent（VIP 模型）")
|
||
p.add_argument("--mask", default=None, help="mask 图路径或 URL")
|
||
p.add_argument(
|
||
"--callback-url",
|
||
default=None,
|
||
help="流程回调 URL：每个事件 POST JSON",
|
||
)
|
||
p.add_argument(
|
||
"--callback-file",
|
||
default=None,
|
||
help="流程回调状态文件（持续覆盖写入 latest/history）",
|
||
)
|
||
p.add_argument("--poll-interval", type=float, default=POLL_INTERVAL_SEC)
|
||
p.add_argument("--poll-timeout", type=float, default=POLL_TIMEOUT_SEC)
|
||
p.add_argument(
|
||
"--json",
|
||
dest="json_input",
|
||
default=None,
|
||
help='JSON 参数（对象字符串 / 文件路径 / "-" 读 stdin）',
|
||
)
|
||
p.add_argument(
|
||
"--result-only",
|
||
action="store_true",
|
||
help="只输出最终结果 JSON（抑制中间 NDJSON 事件到 stdout，仍写回调）",
|
||
)
|
||
p.add_argument(
|
||
"--quiet",
|
||
action="store_true",
|
||
help="不向 stdout 打印事件流（最终结果仍打印）",
|
||
)
|
||
return p
|
||
|
||
|
||
def main(argv: Optional[Sequence[str]] = None) -> int:
|
||
parser = build_parser()
|
||
args = parser.parse_args(argv)
|
||
payload = _load_json_arg(args.json_input)
|
||
|
||
prompt = payload.get("prompt") or args.prompt
|
||
if not prompt:
|
||
parser.error("必须提供 --prompt 或 JSON.prompt")
|
||
|
||
images = payload.get("images") or args.images or []
|
||
if isinstance(images, str):
|
||
images = [images]
|
||
|
||
output = payload.get("output") or args.output
|
||
quiet_events = bool(args.quiet or args.result_only or payload.get("quiet"))
|
||
|
||
try:
|
||
result = generate(
|
||
prompt=prompt,
|
||
output=output,
|
||
images=images,
|
||
model=payload.get("model") or args.model,
|
||
api_key=payload.get("api_key") or args.api_key,
|
||
base_url=payload.get("base_url") or args.base_url,
|
||
aspect_ratio=payload.get("aspectRatio")
|
||
or payload.get("aspect_ratio")
|
||
or args.aspect_ratio,
|
||
quality=payload.get("quality") or args.quality,
|
||
reply_type=payload.get("replyType")
|
||
or payload.get("reply_type")
|
||
or args.reply_type,
|
||
background=payload.get("background") or args.background,
|
||
mask=payload.get("mask") or args.mask,
|
||
callback_url=payload.get("callback_url") or args.callback_url,
|
||
callback_file=payload.get("callback_file") or args.callback_file,
|
||
poll_interval=float(payload.get("poll_interval") or args.poll_interval),
|
||
poll_timeout=float(payload.get("poll_timeout") or args.poll_timeout),
|
||
quiet=quiet_events,
|
||
)
|
||
except GenerateError as exc:
|
||
err = {"ok": False, "error": str(exc)}
|
||
print(json.dumps(err, ensure_ascii=False), flush=True)
|
||
return 1
|
||
except Exception as exc: # noqa: BLE001
|
||
err = {"ok": False, "error": str(exc)}
|
||
print(json.dumps(err, ensure_ascii=False), flush=True)
|
||
return 2
|
||
|
||
# 最终结果单独一行，AI / 流程可取最后一行或 --result-only 模式
|
||
print(json.dumps(result, ensure_ascii=False), flush=True)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|