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 中配置，参考 .env.example。")
|
|||
|
|
if not key:
|
|||
|
|
die("缺少 IMAGE_API_KEY。请在 .env 中配置，参考 .env.example。")
|
|||
|
|
# 允许用户直接填到 /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()
|