187 lines
6.3 KiB
Python
187 lines
6.3 KiB
Python
#!/usr/bin/env python3
|
||
# /// script
|
||
# requires-python = ">=3.10"
|
||
# dependencies = [
|
||
# "requests>=2.31.0",
|
||
# "pillow>=10.0.0",
|
||
# ]
|
||
# ///
|
||
"""使用 Agnes AI（OpenAI 兼容网关）生成图片。
|
||
支持：文生图、图生图（参考图保持一致性）、URL/Base64 输出、下载保存、白底转透明。
|
||
"""
|
||
|
||
import argparse
|
||
import io
|
||
import json
|
||
import os
|
||
import sys
|
||
import time
|
||
import urllib.request
|
||
|
||
import requests
|
||
|
||
API_URL = "https://apihub.agnes-ai.com/v1/images/generations"
|
||
DEFAULT_MODEL = "agnes-image-2.1-flash"
|
||
|
||
|
||
def call_api(prompt, model, size, api_key, images=None, want_b64=False, retries=3):
|
||
"""调用 Agnes 图像生成接口，返回解析后的 JSON。
|
||
对瞬时网络/SSL 抖动自动重试，重试间隔递增。
|
||
"""
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Authorization": f"Bearer {api_key}",
|
||
}
|
||
extra = {"response_format": "b64_json" if want_b64 else "url"}
|
||
if images:
|
||
# 图生图：参考图放在 extra_body.image，支持公网 URL 或 data:image Base64
|
||
extra["image"] = images
|
||
payload = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"size": size,
|
||
"extra_body": extra,
|
||
}
|
||
# 文生图需要 Base64 输出时，按文档使用顶层 return_base64
|
||
if want_b64 and not images:
|
||
payload["return_base64"] = True
|
||
|
||
last_err = None
|
||
for attempt in range(1, retries + 1):
|
||
try:
|
||
resp = requests.post(API_URL, headers=headers, json=payload, timeout=180)
|
||
resp.raise_for_status()
|
||
return resp.json()
|
||
except (requests.exceptions.SSLError,
|
||
requests.exceptions.ConnectionError,
|
||
requests.exceptions.Timeout) as e:
|
||
# 瞬时网络抖动：重试
|
||
last_err = e
|
||
if attempt < retries:
|
||
time.sleep(attempt * 2)
|
||
continue
|
||
raise
|
||
raise last_err
|
||
|
||
|
||
def fetch_bytes(url):
|
||
"""下载 URL 内容为字节。"""
|
||
with urllib.request.urlopen(url, timeout=180) as r:
|
||
return r.read()
|
||
|
||
|
||
def make_transparent(png_bytes, thresh=60, white_min=225):
|
||
"""把纯白背景变透明（保留内部白色）。
|
||
从图像四边多个种子点做 flood fill，只清除与边缘连通的白色区域。
|
||
"""
|
||
from PIL import Image, ImageDraw
|
||
|
||
img = Image.open(io.BytesIO(png_bytes)).convert("RGBA")
|
||
w, h = img.size
|
||
seeds = []
|
||
steps = 12
|
||
for i in range(steps + 1):
|
||
t = int(i * (w - 1) / steps)
|
||
seeds += [(t, 0), (t, h - 1)]
|
||
s = int(i * (h - 1) / steps)
|
||
seeds += [(0, s), (w - 1, s)]
|
||
for sx, sy in seeds:
|
||
r, g, b, a = img.getpixel((sx, sy))
|
||
if a > 0 and r >= white_min and g >= white_min and b >= white_min:
|
||
ImageDraw.floodfill(img, (sx, sy), (255, 255, 255, 0), thresh=thresh)
|
||
out = io.BytesIO()
|
||
img.save(out, "PNG")
|
||
return out.getvalue()
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="使用 Agnes AI 生成图片。")
|
||
parser.add_argument("--prompt", required=True, help="图像生成的文本描述（必选）")
|
||
parser.add_argument("--model", default=DEFAULT_MODEL, help=f"模型 ID，默认 {DEFAULT_MODEL}")
|
||
parser.add_argument("--size", default="1024x1024", help="图像尺寸，如 1024x1024、1024x768")
|
||
parser.add_argument("--image", action="append", default=None,
|
||
help="参考图 URL 或 data:image Base64（图生图，可多次指定）")
|
||
parser.add_argument("--api-key", help="Agnes API Key（也可用 AGNES_API_KEY 环境变量）")
|
||
parser.add_argument("--save", help="下载并保存到本地路径，输出 SAVED:<path>")
|
||
parser.add_argument("--transparent", action="store_true",
|
||
help="配合 --save：把纯白背景转为透明 PNG")
|
||
parser.add_argument("--thresh", type=int, default=60, help="透明化白色容差，默认 60")
|
||
|
||
args = parser.parse_args()
|
||
|
||
api_key = args.api_key or os.environ.get("AGNES_API_KEY")
|
||
if not api_key:
|
||
print("ERROR: 缺少 API Key，请用 --api-key 或设置 AGNES_API_KEY 环境变量。")
|
||
sys.exit(1)
|
||
|
||
# 需要保存且要透明时，直接请求 URL 再下载处理（Pillow 处理本地字节）
|
||
want_b64 = False
|
||
try:
|
||
result = call_api(
|
||
prompt=args.prompt,
|
||
model=args.model,
|
||
size=args.size,
|
||
api_key=api_key,
|
||
images=args.image,
|
||
want_b64=want_b64,
|
||
)
|
||
except requests.exceptions.RequestException as e:
|
||
print(f"ERROR: API 请求失败: {e}")
|
||
if getattr(e, "response", None) is not None:
|
||
print(f"Response body: {e.response.text[:500]}")
|
||
sys.exit(1)
|
||
|
||
data = result.get("data") or []
|
||
if not data:
|
||
print(f"ERROR: 返回中无图像数据。完整响应: {json.dumps(result)[:500]}")
|
||
sys.exit(1)
|
||
|
||
item = data[0]
|
||
url = item.get("url")
|
||
b64 = item.get("b64_json")
|
||
|
||
if not args.save:
|
||
# 仅输出链接（对齐 seedream 约定）
|
||
if url:
|
||
print(f"MEDIA_URL: {url}")
|
||
elif b64:
|
||
print("ERROR: 收到 Base64 但未指定 --save，无法直接展示。请加 --save 保存。")
|
||
sys.exit(1)
|
||
else:
|
||
print(f"ERROR: 无 url 也无 b64_json。响应: {json.dumps(result)[:500]}")
|
||
sys.exit(1)
|
||
return
|
||
|
||
# 保存到本地
|
||
try:
|
||
if b64:
|
||
import base64
|
||
img_bytes = base64.b64decode(b64)
|
||
elif url:
|
||
img_bytes = fetch_bytes(url)
|
||
else:
|
||
print("ERROR: 无可保存的图像数据。")
|
||
sys.exit(1)
|
||
|
||
if args.transparent:
|
||
try:
|
||
img_bytes = make_transparent(img_bytes, thresh=args.thresh)
|
||
except ImportError:
|
||
print("ERROR: --transparent 需要 Pillow，请先 `pip install pillow`。")
|
||
sys.exit(1)
|
||
|
||
os.makedirs(os.path.dirname(os.path.abspath(args.save)), exist_ok=True)
|
||
with open(args.save, "wb") as f:
|
||
f.write(img_bytes)
|
||
# 即使保存到本地，也输出原始图片链接（便于直接引用/展示）
|
||
if url:
|
||
print(f"MEDIA_URL: {url}")
|
||
print(f"SAVED: {args.save}")
|
||
except Exception as e:
|
||
print(f"ERROR: 保存失败: {e}")
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|