137 lines
4.9 KiB
Python
137 lines
4.9 KiB
Python
# LIB · Seedream 3/4/5 适配器
|
||
# D181 · 2026-07-10 · 默认模型: Seedream 3.0 (轻量,便宜)
|
||
"""火山方舟 ARK API → Seedream 系列 文生图/图生图"""
|
||
import json
|
||
import os
|
||
import urllib.request
|
||
import base64
|
||
|
||
from lib.secrets import get_ark_key
|
||
ARK_KEY = get_ark_key()
|
||
ARK_BASE = "https://ark.cn-beijing.volces.com/api/v3"
|
||
|
||
# 模型配置
|
||
# 模型配置（ARK端点上实际可用的模型）
|
||
SEEDREAM_3 = "doubao-seedream-4-0-250828" # Seedream 4.0 (最便宜可用)
|
||
SEEDREAM_4 = "doubao-seedream-4-0-250828" # 同上
|
||
SEEDREAM_4_5 = "doubao-seedream-4-5-251128" # Seedream 4.5
|
||
SEEDREAM_5 = "doubao-seedream-5-0-260128" # Seedream 5.0
|
||
SEEDREAM_5_PRO = "doubao-seedream-5-0-pro-260628" # Seedream 5.0 Pro
|
||
SEEDREAM_MODEL = SEEDREAM_3 # 默认模型
|
||
SEEDREAM_MODEL = SEEDREAM_3 # 默认模型 (爸爸指定用便宜的)
|
||
|
||
# 尺寸约束
|
||
SIZES = {
|
||
SEEDREAM_3: ["1024x1024", "1024x768", "768x1024"],
|
||
SEEDREAM_4: ["1664x936", "936x1664", "1024x1024"],
|
||
}
|
||
DEFAULT_SIZE = "1024x1024" # 即梦3默认方图
|
||
|
||
def _api_call(endpoint, payload):
|
||
"""通用 ARK API 调用"""
|
||
req = urllib.request.Request(
|
||
f"{ARK_BASE}{endpoint}",
|
||
data=json.dumps(payload).encode(),
|
||
headers={
|
||
"Authorization": f"Bearer {ARK_KEY}",
|
||
"Content-Type": "application/json"
|
||
},
|
||
method="POST"
|
||
)
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||
return json.loads(resp.read())
|
||
except urllib.error.HTTPError as e:
|
||
err_body = e.read().decode()
|
||
return {"error": str(e.code), "detail": err_body}
|
||
|
||
def seedream_t2i(prompt, size=DEFAULT_SIZE, n=1, model=SEEDREAM_MODEL):
|
||
"""文生图
|
||
prompt: 中英文提示词
|
||
size: "1024x1024" (即梦3) / "1664x936" (即梦4)
|
||
model: SEEDREAM_3 (默认) / SEEDREAM_4
|
||
"""
|
||
# 即梦3不支持 n>1
|
||
if model == SEEDREAM_3 and n > 1:
|
||
import warnings
|
||
warnings.warn("即梦3.0 不支持 n>1，已固定为 n=1")
|
||
n = 1
|
||
payload = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"size": size,
|
||
"n": n
|
||
}
|
||
return _api_call("/images/generations", payload)
|
||
|
||
def seedream_i2i(prompt, image_path, size=DEFAULT_SIZE, n=1, model=SEEDREAM_MODEL):
|
||
"""图生图 · 即梦3.0 图生图 (seededit-3.0)"""
|
||
if model == SEEDREAM_3:
|
||
# 即梦3图生图用 seededit-3.0
|
||
model = "seededit-3.0"
|
||
with open(image_path, "rb") as f:
|
||
img_b64 = base64.b64encode(f.read()).decode()
|
||
payload = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"image": f"data:image/jpeg;base64,{img_b64}",
|
||
"size": size,
|
||
"n": 1 # 图生图只支持单张
|
||
}
|
||
return _api_call("/images/generations", payload)
|
||
|
||
def seedream_multi(prompt, image_paths, size=DEFAULT_SIZE, n=1, model=SEEDREAM_MODEL):
|
||
"""多图融合 · 仅即梦4支持"""
|
||
if model == SEEDREAM_3:
|
||
raise RuntimeError("多图融合仅即梦4.0支持，请使用 model=SEEDREAM_4")
|
||
imgs = []
|
||
for p in image_paths:
|
||
with open(p, "rb") as f:
|
||
fmt = "png" if p.endswith(".png") else "jpeg"
|
||
imgs.append(f"data:image/{fmt};base64,{base64.b64encode(f.read()).decode()}")
|
||
payload = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"image": imgs,
|
||
"size": size,
|
||
"n": n
|
||
}
|
||
return _api_call("/images/generations", payload)
|
||
|
||
def seedream_download(result, output_path):
|
||
"""从 API 结果下载图片"""
|
||
if "error" in result:
|
||
return None
|
||
url = result.get("data", [{}])[0].get("url")
|
||
if not url:
|
||
return None
|
||
urllib.request.urlretrieve(url, output_path)
|
||
return output_path
|
||
|
||
# ====== CLI ======
|
||
if __name__ == "__main__":
|
||
import sys
|
||
if len(sys.argv) < 2:
|
||
print("Usage: python seedream.py t2i <prompt> [output] [--model 3|4] [--size WxH]")
|
||
print(" python seedream.py i2i <prompt> <image_path> [output] [--model 3|4]")
|
||
print(" python seedream.py multi <prompt> <img1> <img2> [output]")
|
||
print(f" Default model: {SEEDREAM_MODEL} (即梦3.0 · ¥0.2/张)")
|
||
sys.exit(1)
|
||
cmd = sys.argv[1]
|
||
if cmd == "t2i":
|
||
prompt = sys.argv[2]
|
||
out = sys.argv[3] if len(sys.argv) > 3 else "seedream_output.jpg"
|
||
r = seedream_t2i(prompt)
|
||
p = seedream_download(r, out)
|
||
print(f"✅ {p}" if p else f"❌ {r.get('error','unknown')}")
|
||
elif cmd == "i2i":
|
||
prompt, img, out = sys.argv[2], sys.argv[3], (sys.argv[4] if len(sys.argv) > 4 else "seedream_i2i_output.jpg")
|
||
r = seedream_i2i(prompt, img)
|
||
p = seedream_download(r, out)
|
||
print(f"✅ {p}" if p else f"❌ {r.get('error','unknown')}")
|
||
elif cmd == "multi":
|
||
prompt, imgs, out = sys.argv[2], sys.argv[3:-1], sys.argv[-1]
|
||
r = seedream_multi(prompt, imgs)
|
||
p = seedream_download(r, out)
|
||
print(f"✅ {p}" if p else f"❌ {r.get('error','unknown')}")
|