150 lines
5.1 KiB
Python
Executable File
150 lines
5.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""GLM-Image 文生图调用脚本（智谱 GLM-Image 模型）。
|
||
|
||
用法:
|
||
python3 generate.py "提示词" [-s SIZE] [-o OUTPUT] [--open] [--json]
|
||
|
||
示例:
|
||
python3 generate.py "一只可爱的小猫咪，坐在阳光明媚的窗台上" -s 1280x1280
|
||
python3 generate.py "商业海报：新品上市" -s 1056x1568 -o poster.png --open
|
||
|
||
默认 size=1280x1280，输出到当前目录 glm-image-<时间戳>.png
|
||
API Key 优先读环境变量 GLM_API_KEY，否则用内置默认 key。
|
||
"""
|
||
import argparse
|
||
import json
|
||
import os
|
||
import sys
|
||
import time
|
||
import urllib.request
|
||
import urllib.error
|
||
|
||
API_ENDPOINT = "https://open.bigmodel.cn/api/paas/v4/images/generations"
|
||
DEFAULT_KEY = "" # 服务器版不含内置 key，请设置环境变量 GLM_API_KEY
|
||
|
||
RECOMMENDED_SIZES = [
|
||
"1280x1280", "1568x1056", "1056x1568",
|
||
"1472x1088", "1088x1472", "1728x960", "960x1728",
|
||
]
|
||
|
||
|
||
def parse_size(size):
|
||
"""校验 size，返回 (w, h)。规则：512-2048，且为 32 的整数倍。"""
|
||
try:
|
||
w, h = size.lower().split("x")
|
||
w, h = int(w), int(h)
|
||
except ValueError:
|
||
raise ValueError(f"size 格式错误：'{size}'，应为 WxH，如 1280x1280")
|
||
for v, name in ((w, "宽"), (h, "高")):
|
||
if v < 512 or v > 2048:
|
||
raise ValueError(f"{name}={v} 不在 512-2048 范围内")
|
||
if v % 32 != 0:
|
||
raise ValueError(f"{name}={v} 不是 32 的整数倍")
|
||
return w, h
|
||
|
||
|
||
def generate(prompt, size="1280x1280", api_key=None, timeout=120, watermark=True):
|
||
"""调用 GLM-Image 接口，返回图片 URL。
|
||
|
||
watermark=False 关闭 AI 水印，需账号已在「个人中心-安全管理-去水印管理」签署免责声明。
|
||
"""
|
||
parse_size(size) # 校验
|
||
key = api_key or os.environ.get("GLM_API_KEY") or DEFAULT_KEY
|
||
payload = {
|
||
"model": "glm-image",
|
||
"prompt": prompt,
|
||
"size": size,
|
||
"watermark_enabled": watermark,
|
||
}
|
||
data = json.dumps(payload).encode("utf-8")
|
||
req = urllib.request.Request(
|
||
API_ENDPOINT,
|
||
data=data,
|
||
headers={
|
||
"Authorization": f"Bearer {key}",
|
||
"Content-Type": "application/json",
|
||
"Accept": "application/json",
|
||
},
|
||
method="POST",
|
||
)
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
body = resp.read().decode("utf-8")
|
||
except urllib.error.HTTPError as e:
|
||
err = e.read().decode("utf-8", errors="replace")
|
||
raise RuntimeError(f"HTTP {e.code}: {err}") from None
|
||
except urllib.error.URLError as e:
|
||
raise RuntimeError(f"网络错误: {e.reason}") from None
|
||
|
||
obj = json.loads(body)
|
||
if not obj.get("data"):
|
||
raise RuntimeError(f"返回无 data 字段: {body}")
|
||
url = obj["data"][0].get("url")
|
||
if not url:
|
||
raise RuntimeError(f"返回无 url: {body}")
|
||
return url, obj
|
||
|
||
|
||
def download(url, output):
|
||
"""下载图片 URL 到 output，返回输出路径。"""
|
||
with urllib.request.urlopen(url, timeout=120) as resp:
|
||
content = resp.read()
|
||
with open(output, "wb") as f:
|
||
f.write(content)
|
||
return output
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser(description="GLM-Image 文生图")
|
||
ap.add_argument("prompt", help="生成提示词（最多 1000 字符）")
|
||
ap.add_argument("-s", "--size", default="1280x1280",
|
||
help=f"图片尺寸 WxH（默认 1280x1280）。推荐: {', '.join(RECOMMENDED_SIZES)}")
|
||
ap.add_argument("-o", "--output", default=None,
|
||
help="输出文件路径（默认 glm-image-<时间戳>.png）")
|
||
ap.add_argument("--open", action="store_true", help="生成后在 Finder 中打开")
|
||
ap.add_argument("--json", action="store_true", help="打印完整 JSON 返回")
|
||
ap.add_argument("--no-download", action="store_true", help="只返回 URL，不下载")
|
||
ap.add_argument("--no-watermark", action="store_true",
|
||
help="关闭 AI 水印（需账号已在「个人中心-安全管理-去水印管理」签署免责声明）")
|
||
args = ap.parse_args()
|
||
|
||
if len(args.prompt) > 1000:
|
||
sys.exit(f"错误: 提示词 {len(args.prompt)} 字符，超过 1000 上限")
|
||
|
||
try:
|
||
parse_size(args.size)
|
||
except ValueError as e:
|
||
sys.exit(f"错误: {e}")
|
||
|
||
print(f"→ 调用 GLM-Image（size={args.size}, watermark={'off' if args.no_watermark else 'on'}）...", file=sys.stderr)
|
||
t0 = time.time()
|
||
try:
|
||
url, obj = generate(args.prompt, args.size, watermark=not args.no_watermark)
|
||
except RuntimeError as e:
|
||
sys.exit(f"错误: {e}")
|
||
elapsed = time.time() - t0
|
||
print(f"✓ 生成成功（{elapsed:.1f}s）: {url}", file=sys.stderr)
|
||
|
||
if args.json:
|
||
print(json.dumps(obj, ensure_ascii=False, indent=2))
|
||
|
||
if args.no_download:
|
||
print(url)
|
||
return
|
||
|
||
output = args.output or f"glm-image-{int(time.time())}.png"
|
||
try:
|
||
download(url, output)
|
||
except Exception as e:
|
||
sys.exit(f"下载失败: {e}\n图片 URL: {url}")
|
||
print(f"✓ 已保存: {output}", file=sys.stderr)
|
||
print(output)
|
||
|
||
if args.open:
|
||
os.system(f'open "{output}"')
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|