小楼
221 lines
8.7 KiB
Python
221 lines
8.7 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
gen_image.py — 通过大模型中转站调用 OpenAI 图片接口生图（文生图 / 图生图）。
|
||
|
||
零第三方依赖，仅用 Python 标准库。凭证从环境变量或 skill 目录下的 .env 读取。
|
||
|
||
用法:
|
||
# 文生图
|
||
python3 gen_image.py "一只戴墨镜的柴犬，赛博朋克霓虹背景" --size 1024x1024 --n 1
|
||
|
||
# 图生图 / 编辑（可选 --mask 做局部重绘）
|
||
python3 gen_image.py "把背景换成雪山" --edit input.png --mask mask.png
|
||
|
||
环境变量（也可写进 .env）:
|
||
IMAGE_API_BASE_URL 中转站基址，如 https://your-relay.com (会自动补 /v1/images/...)
|
||
IMAGE_API_KEY API Key
|
||
IMAGE_MODEL 默认模型，如 gpt-image-2（默认 gpt-image-2）
|
||
"""
|
||
|
||
import argparse
|
||
import base64
|
||
import json
|
||
import mimetypes
|
||
import os
|
||
import sys
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
import uuid
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 配置加载
|
||
# ---------------------------------------------------------------------------
|
||
def load_dotenv():
|
||
"""从脚本上级目录（skill 根）的 .env 读取 KEY=VALUE，不覆盖已存在的环境变量。"""
|
||
skill_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
env_path = os.path.join(skill_root, ".env")
|
||
if not os.path.isfile(env_path):
|
||
return
|
||
with open(env_path, "r", encoding="utf-8") as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if not line or line.startswith("#") or "=" not in line:
|
||
continue
|
||
key, _, val = line.partition("=")
|
||
key, val = key.strip(), val.strip().strip('"').strip("'")
|
||
os.environ.setdefault(key, val)
|
||
|
||
|
||
def get_config():
|
||
base = os.environ.get("IMAGE_API_BASE_URL", "").rstrip("/")
|
||
key = os.environ.get("IMAGE_API_KEY", "")
|
||
model = os.environ.get("IMAGE_MODEL", "gpt-image-2")
|
||
if not base:
|
||
die("缺少 IMAGE_API_BASE_URL（中转站基址）。请在工具箱「中转」里勾选用于生图，或在 .env 中配置。")
|
||
if not key:
|
||
die("缺少 IMAGE_API_KEY。请在工具箱「中转」里勾选用于生图，或在 .env 中配置。")
|
||
# 允许用户直接填到 /v1，做个兼容归一
|
||
if base.endswith("/v1"):
|
||
base = base[: -len("/v1")]
|
||
return base, key, model
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# HTTP 辅助
|
||
# ---------------------------------------------------------------------------
|
||
def die(msg, code=1):
|
||
print(f"[gen_image] 错误: {msg}", file=sys.stderr)
|
||
sys.exit(code)
|
||
|
||
|
||
# 部分中转站前置 Cloudflare/WAF，会按 UA 拦截默认的 Python-urllib（error 1010）。
|
||
# 用常规浏览器 UA 规避。
|
||
BROWSER_UA = (
|
||
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
|
||
"AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36"
|
||
)
|
||
|
||
|
||
def _request(url, key, data, content_type):
|
||
req = urllib.request.Request(url, data=data, method="POST")
|
||
req.add_header("Authorization", f"Bearer {key}")
|
||
req.add_header("Content-Type", content_type)
|
||
req.add_header("User-Agent", BROWSER_UA)
|
||
req.add_header("Accept", "application/json")
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=300) as resp:
|
||
return json.loads(resp.read().decode("utf-8"))
|
||
except urllib.error.HTTPError as e:
|
||
body = e.read().decode("utf-8", "replace")
|
||
die(f"HTTP {e.code} — {body[:800]}")
|
||
except urllib.error.URLError as e:
|
||
die(f"网络错误: {e.reason}")
|
||
|
||
|
||
def post_json(url, key, payload):
|
||
return _request(url, key, json.dumps(payload).encode("utf-8"), "application/json")
|
||
|
||
|
||
def post_multipart(url, key, fields, files):
|
||
"""fields: dict[str,str]; files: list[(field, filepath)]。手写 multipart，无依赖。"""
|
||
boundary = "----genimg" + uuid.uuid4().hex
|
||
buf = bytearray()
|
||
|
||
def w(s):
|
||
buf.extend(s.encode("utf-8") if isinstance(s, str) else s)
|
||
|
||
for name, value in fields.items():
|
||
if value is None:
|
||
continue
|
||
w(f"--{boundary}\r\n")
|
||
w(f'Content-Disposition: form-data; name="{name}"\r\n\r\n')
|
||
w(f"{value}\r\n")
|
||
|
||
for name, path in files:
|
||
if not path:
|
||
continue
|
||
fname = os.path.basename(path)
|
||
ctype = mimetypes.guess_type(path)[0] or "application/octet-stream"
|
||
with open(path, "rb") as fh:
|
||
content = fh.read()
|
||
w(f"--{boundary}\r\n")
|
||
w(f'Content-Disposition: form-data; name="{name}"; filename="{fname}"\r\n')
|
||
w(f"Content-Type: {ctype}\r\n\r\n")
|
||
w(content)
|
||
w("\r\n")
|
||
|
||
w(f"--{boundary}--\r\n")
|
||
return _request(url, key, bytes(buf), f"multipart/form-data; boundary={boundary}")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 结果保存
|
||
# ---------------------------------------------------------------------------
|
||
def save_images(resp, out_dir, prefix):
|
||
data = resp.get("data")
|
||
if not data:
|
||
die(f"响应中没有图片数据。原始响应: {json.dumps(resp)[:800]}")
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
saved = []
|
||
for i, item in enumerate(data):
|
||
ts = time.strftime("%Y%m%d-%H%M%S")
|
||
fname = f"{prefix}-{ts}-{i}.png"
|
||
fpath = os.path.join(out_dir, fname)
|
||
if item.get("b64_json"):
|
||
with open(fpath, "wb") as f:
|
||
f.write(base64.b64decode(item["b64_json"]))
|
||
elif item.get("url"):
|
||
try:
|
||
dl = urllib.request.Request(item["url"], headers={"User-Agent": BROWSER_UA})
|
||
with urllib.request.urlopen(dl, timeout=120) as r:
|
||
with open(fpath, "wb") as f:
|
||
f.write(r.read())
|
||
except Exception as e:
|
||
die(f"下载图片失败 {item['url']}: {e}")
|
||
else:
|
||
die(f"未知的图片条目格式: {json.dumps(item)[:400]}")
|
||
saved.append(fpath)
|
||
return saved
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 主流程
|
||
# ---------------------------------------------------------------------------
|
||
def main():
|
||
load_dotenv()
|
||
p = argparse.ArgumentParser(description="通过中转站调用 OpenAI 图片接口生图")
|
||
p.add_argument("prompt", help="生图提示词")
|
||
p.add_argument("--model", default=None, help="覆盖默认模型")
|
||
p.add_argument("--size", default="1024x1024", help="尺寸，如 1024x1024 / 1536x1024 / 1024x1536 / auto")
|
||
p.add_argument("--n", type=int, default=1, help="生成张数")
|
||
p.add_argument("--quality", default=None, help="质量: low/medium/high/auto（gpt-image 系）")
|
||
p.add_argument("--background", default=None, help="transparent/opaque/auto（gpt-image 系，透明背景）")
|
||
p.add_argument("--edit", action="append", default=[], help="图生图/编辑的输入图（可多次传入多图）")
|
||
p.add_argument("--mask", default=None, help="局部重绘蒙版 PNG（透明处为要重绘的区域）")
|
||
p.add_argument("--output", default=None, help="输出目录（默认 skill 根下 output/）")
|
||
args = p.parse_args()
|
||
|
||
base, key, default_model = get_config()
|
||
model = args.model or default_model
|
||
skill_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
out_dir = args.output or os.path.join(skill_root, "output")
|
||
|
||
if args.edit:
|
||
# 图生图 / 编辑：multipart 到 /v1/images/edits
|
||
url = f"{base}/v1/images/edits"
|
||
fields = {"model": model, "prompt": args.prompt, "n": str(args.n), "size": args.size}
|
||
if args.quality:
|
||
fields["quality"] = args.quality
|
||
if args.background:
|
||
fields["background"] = args.background
|
||
files = [("image[]", pth) for pth in args.edit] # 多图统一用 image[]
|
||
if len(args.edit) == 1:
|
||
files = [("image", args.edit[0])] # 单图用 image，兼容 dall-e-2 风格
|
||
if args.mask:
|
||
files.append(("mask", args.mask))
|
||
print(f"[gen_image] 图生图 · model={model} · {len(args.edit)} 张输入 · size={args.size}")
|
||
resp = post_multipart(url, key, fields, files)
|
||
prefix = "edit"
|
||
else:
|
||
# 文生图：JSON 到 /v1/images/generations
|
||
url = f"{base}/v1/images/generations"
|
||
payload = {"model": model, "prompt": args.prompt, "n": args.n, "size": args.size}
|
||
if args.quality:
|
||
payload["quality"] = args.quality
|
||
if args.background:
|
||
payload["background"] = args.background
|
||
print(f"[gen_image] 文生图 · model={model} · n={args.n} · size={args.size}")
|
||
resp = post_json(url, key, payload)
|
||
prefix = "img"
|
||
|
||
saved = save_images(resp, out_dir, prefix)
|
||
print(f"[gen_image] 完成，已保存 {len(saved)} 张:")
|
||
for s in saved:
|
||
print(f" {s}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|