282 lines
12 KiB
Python
282 lines
12 KiB
Python
# LIB · Seedance 2.0 视频生成适配器 + 异步任务队列
|
||
# D180 · 2026-07-10
|
||
"""火山方舟 ARK → Seedance 2.0 · 提交+轮询+下载+写库 全自动"""
|
||
import json, os, subprocess, time, tempfile, base64
|
||
|
||
from lib.secrets import get_ark_key
|
||
ARK_KEY = get_ark_key()
|
||
ARK_BASE = "https://ark.cn-beijing.volces.com/api/v3"
|
||
SEEDANCE_MODEL = "doubao-seedance-2-0-260128"
|
||
TASK_URL = f"{ARK_BASE}/contents/generations/tasks"
|
||
|
||
# 成本估算（元/秒·720p）· 参考价，实际以火山方舟计费为准
|
||
COST_PER_SECOND = {"720p": 0.5, "1080p": 1.0}
|
||
|
||
def _curl(method, url, payload=None):
|
||
"""通用 curl 调用"""
|
||
tf = None
|
||
cmd = ["curl", "-s", url, "-H", f"Authorization: Bearer {ARK_KEY}"]
|
||
if payload:
|
||
tf = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, encoding='utf-8')
|
||
json.dump(payload, tf, ensure_ascii=False)
|
||
tf.close()
|
||
cmd += ["-H", "Content-Type: application/json", "-d", f"@{tf.name}"]
|
||
if method == "POST":
|
||
cmd.insert(1, "-X")
|
||
cmd.insert(2, "POST")
|
||
try:
|
||
r = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
|
||
if tf: os.unlink(tf.name)
|
||
return json.loads(r.stdout) if r.stdout else {"error": r.stderr}
|
||
except Exception as e:
|
||
if tf: os.unlink(tf.name)
|
||
return {"error": str(e)}
|
||
|
||
def submit_task(prompt, references=None, duration=6, resolution="720p"):
|
||
"""提交 Seedance 生成任务 → 返回 task_id"""
|
||
content = [{"type": "text", "text": prompt}]
|
||
if references:
|
||
for ref_type, ref_path in references:
|
||
with open(ref_path, "rb") as f:
|
||
fmt = "png" if ref_path.endswith(".png") else "jpeg"
|
||
b64 = base64.b64encode(f.read()).decode()
|
||
content.append({
|
||
"type": "image_url",
|
||
"role": ref_type,
|
||
"image_url": {"url": f"data:image/{fmt};base64,{b64}"}
|
||
})
|
||
payload = {
|
||
"model": SEEDANCE_MODEL,
|
||
"content": content,
|
||
"duration": duration,
|
||
"resolution": resolution
|
||
}
|
||
result = _curl("POST", TASK_URL, payload)
|
||
if "id" in result:
|
||
return result["id"]
|
||
return {"error": result}
|
||
|
||
def poll_task(task_id, poll_interval=30, max_wait=600):
|
||
"""轮询任务状态 → 返回结果"""
|
||
elapsed = 0
|
||
while elapsed < max_wait:
|
||
result = _curl("GET", f"{TASK_URL}/{task_id}")
|
||
if "error" in result:
|
||
return result
|
||
status = result.get("status", "")
|
||
if status == "succeeded":
|
||
return result
|
||
if status == "failed":
|
||
return {"error": f"Task failed: {result}"}
|
||
time.sleep(poll_interval)
|
||
elapsed += poll_interval
|
||
return {"error": "timeout"}
|
||
|
||
def download_video(video_url, output_path):
|
||
"""下载视频"""
|
||
subprocess.run(["curl", "-s", "-o", output_path, video_url], check=False)
|
||
return os.path.exists(output_path) and os.path.getsize(output_path) > 0
|
||
|
||
def generate_shot(shot_id, prompt, references, duration=6, output_dir="outputs/videos"):
|
||
"""一站式生成：提交→轮询→下载→返回结果"""
|
||
print(f"[LIB] 提交 {shot_id}...")
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
task_id = submit_task(prompt, references, duration)
|
||
if isinstance(task_id, dict):
|
||
return {"error": task_id}
|
||
|
||
print(f"[LIB] 任务 {task_id[:20]}... 等待完成")
|
||
result = poll_task(task_id)
|
||
if "error" in result:
|
||
return result
|
||
|
||
video_url = result.get("content", {}).get("video_url")
|
||
if not video_url:
|
||
return {"error": "no video_url"}
|
||
|
||
output_path = os.path.join(output_dir, f"{shot_id}.mp4")
|
||
print(f"[LIB] 下载视频...")
|
||
if download_video(video_url, output_path):
|
||
return {
|
||
"task_id": task_id,
|
||
"output_path": output_path,
|
||
"seed": result.get("seed"),
|
||
"duration": result.get("duration"),
|
||
"resolution": result.get("resolution"),
|
||
"tokens": result.get("usage", {}).get("total_tokens", 0)
|
||
}
|
||
return {"error": "download failed"}
|
||
|
||
def generate_shot_async(shot_id, prompt, references, duration=6, output_dir="outputs/videos"):
|
||
"""异步提交：只提交不等待，返回 task_id 用于后续轮询"""
|
||
task_id = submit_task(prompt, references, duration)
|
||
if isinstance(task_id, dict):
|
||
return task_id
|
||
# 写入任务记录（含成本估算）
|
||
from lib.models import get_db
|
||
db = get_db()
|
||
estimated_cost = round(duration * COST_PER_SECOND.get("720p", 0.5), 2)
|
||
db.execute("""
|
||
INSERT INTO generation_tasks (id, shot_id, task_type, seedance_task_id, status, prompt_used, references_used, cost_estimated, started_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, datetime('now'))
|
||
""", (task_id[:20], shot_id, "video", task_id, "running", prompt, json.dumps([r[0] for r in references] if references else []), estimated_cost))
|
||
db.execute("UPDATE shots SET seedance_task_id = ?, status = 'generating' WHERE id = ?", (task_id, shot_id))
|
||
db.commit()
|
||
return {"task_id": task_id, "status": "submitted"}
|
||
|
||
def check_and_collect(shot_id):
|
||
"""检查任务状态，完成后自动下载"""
|
||
from lib.models import get_db
|
||
db = get_db()
|
||
task = db.execute("""
|
||
SELECT * FROM generation_tasks WHERE shot_id = ? ORDER BY created_at DESC LIMIT 1
|
||
""", (shot_id,)).fetchone()
|
||
if not task:
|
||
return {"error": "no task found"}
|
||
result = poll_task(task["seedance_task_id"], poll_interval=0, max_wait=1)
|
||
if "error" in result:
|
||
return result
|
||
status = result.get("status", "")
|
||
if status in ("running", "queued", "processing"):
|
||
return {"status": "running", "task_id": task["seedance_task_id"]}
|
||
if status != "succeeded":
|
||
return {"error": f"Task {status}: {str(result)[:200]}"}
|
||
|
||
# Succeeded - download
|
||
video_url = result.get("content", {}).get("video_url")
|
||
output_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "outputs", "videos")
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
output_path = os.path.join(output_dir, f"{shot_id}.mp4")
|
||
if download_video(video_url, output_path):
|
||
db.execute("""
|
||
UPDATE generation_tasks SET status = 'succeeded', cost_actual = ?, completed_at = datetime('now') WHERE id = ?
|
||
""", (result.get("usage", {}).get("total_tokens", 0), task["id"],))
|
||
# 保存seed
|
||
seed = result.get("seed")
|
||
if seed:
|
||
db.execute("UPDATE shots SET status = 'done', output_path = ?, seed_value = ? WHERE id = ?",
|
||
(output_path, seed, shot_id))
|
||
else:
|
||
db.execute("UPDATE shots SET status = 'done', output_path = ? WHERE id = ?",
|
||
(output_path, shot_id))
|
||
db.commit()
|
||
return {"status": "done", "output_path": output_path, "seed": result.get("seed")}
|
||
return {"error": "download failed"}
|
||
|
||
# ====== 批量并行生成 + 自动重试 ======
|
||
|
||
def batch_submit(episode_id, max_parallel=3, retry_failed=False):
|
||
"""批量提交：自动读库中所有 pending 镜头 → 并行提交 → 并行轮询
|
||
|
||
支持 Ctrl+C 中断恢复：已提交的任务 ID 都在数据库中，
|
||
中断后用 batch_resume() 继续收集。
|
||
"""
|
||
from lib.models import get_db
|
||
db = get_db()
|
||
|
||
where = "episode_id = ? AND status IN ('pending', 'failed')" if retry_failed else "episode_id = ? AND status = 'pending'"
|
||
shots = db.execute(f"SELECT * FROM shots WHERE {where} ORDER BY shot_number", (episode_id,)).fetchall()
|
||
|
||
if not shots:
|
||
print("⚠️ 没有待生成的镜头")
|
||
return
|
||
|
||
print(f"🚀 批量提交 {len(shots)} 个镜头（最多并行 {max_parallel}）")
|
||
print(f"💡 可随时 Ctrl+C 中断，用 batch-resume 恢复\n")
|
||
|
||
results = {}
|
||
queue = list(shots)
|
||
|
||
while queue or any(r.get("status") == "running" for r in results.values()):
|
||
while len([r for r in results.values() if r.get("status") in ("submitted", "running")]) < max_parallel and queue:
|
||
shot = queue.pop(0)
|
||
shot_id = shot["id"]
|
||
print(f"📤 提交 {shot_id}...")
|
||
refs = db.execute("SELECT sa.role, a.file_path FROM shot_assets sa JOIN assets a ON sa.asset_id = a.id WHERE sa.shot_id = ?", (shot_id,)).fetchall()
|
||
ref_paths = [(r["role"], r["file_path"]) for r in refs] if refs else []
|
||
prompt = shot["prompt"] or shot["description"]
|
||
r = generate_shot_async(shot_id, prompt, ref_paths, int(shot["duration_sec"] or 6))
|
||
if "error" in r:
|
||
results[shot_id] = {"status": "failed", "error": r["error"], "retries": 0}
|
||
print(f" ❌ {r['error']}")
|
||
else:
|
||
results[shot_id] = {"status": "running", "task_id": r["task_id"], "retries": 0}
|
||
print(f" ✅ {r['task_id'][:20]}...")
|
||
|
||
for shot_id, info in list(results.items()):
|
||
if info["status"] != "running": continue
|
||
result = check_and_collect(shot_id)
|
||
if result.get("status") == "done":
|
||
results[shot_id] = {"status": "done", "output_path": result["output_path"]}
|
||
print(f"✅ {shot_id} 完成")
|
||
elif result.get("error") and "timeout" not in str(result.get("error", "")):
|
||
info["retries"] += 1
|
||
if info["retries"] <= 2:
|
||
print(f"🔄 {shot_id} 重试 ({info['retries']}/2)")
|
||
db.execute("UPDATE shots SET status = 'pending' WHERE id = ?", (shot_id,)); db.commit()
|
||
refs = db.execute("SELECT sa.role, a.file_path FROM shot_assets sa JOIN assets a ON sa.asset_id = a.id WHERE sa.shot_id = ?", (shot_id,)).fetchall()
|
||
ref_paths = [(r["role"], r["file_path"]) for r in refs]
|
||
shot = db.execute("SELECT * FROM shots WHERE id = ?", (shot_id,)).fetchone()
|
||
prompt = shot["prompt"] or shot["description"]
|
||
new_r = generate_shot_async(shot_id, prompt, ref_paths, int(shot["duration_sec"] or 6))
|
||
if "error" not in new_r:
|
||
results[shot_id] = {"status": "running", "task_id": new_r["task_id"], "retries": info["retries"]}
|
||
else:
|
||
results[shot_id] = {"status": "failed", "error": new_r["error"], "retries": info["retries"]}
|
||
else:
|
||
results[shot_id] = {"status": "failed", "retries": info["retries"]}
|
||
db.execute("UPDATE shots SET status = 'failed', error_log = ? WHERE id = ?", (str(result.get("error", ""))[:200], shot_id)); db.commit()
|
||
print(f"❌ {shot_id} 重试{info['retries']}次后失败")
|
||
time.sleep(30)
|
||
|
||
done = sum(1 for r in results.values() if r["status"] == "done")
|
||
failed = sum(1 for r in results.values() if r["status"] == "failed")
|
||
print(f"\n📊 批量完成: {done}✅ {failed}❌ (共{len(results)})")
|
||
return results
|
||
|
||
|
||
def batch_resume(episode_id):
|
||
"""从中断恢复：读数据库中所有 generating 状态的镜头 → 继续轮询收集
|
||
|
||
适用于 batch_submit 中途 Ctrl+C 后恢复。
|
||
已提交到 Seedance 的任务 ID 保存在 generation_tasks 表中，不会丢失。
|
||
"""
|
||
from lib.models import get_db
|
||
db = get_db()
|
||
|
||
tasks = db.execute("""
|
||
SELECT gt.shot_id, gt.seedance_task_id, s.shot_number
|
||
FROM generation_tasks gt
|
||
JOIN shots s ON gt.shot_id = s.id
|
||
WHERE s.episode_id = ? AND s.status = 'generating'
|
||
ORDER BY s.shot_number
|
||
""", (episode_id,)).fetchall()
|
||
|
||
if not tasks:
|
||
print("⚠️ 没有正在生成的镜头，无需恢复")
|
||
return
|
||
|
||
print(f"🔄 恢复 {len(tasks)} 个正在生成的镜头...\n")
|
||
|
||
results = {}
|
||
for task in tasks:
|
||
shot_id = task["shot_id"]
|
||
print(f"📥 检查 {shot_id} (task: {task['seedance_task_id'][:20]}...)")
|
||
result = check_and_collect(shot_id)
|
||
if result.get("status") == "done":
|
||
results[shot_id] = "done"
|
||
print(f" ✅ 已完成")
|
||
elif result.get("status") == "running":
|
||
results[shot_id] = "running"
|
||
print(f" 🔄 仍在生成中")
|
||
else:
|
||
results[shot_id] = "unknown"
|
||
print(f" ❓ {result}")
|
||
|
||
done = sum(1 for v in results.values() if v == "done")
|
||
running = sum(1 for v in results.values() if v == "running")
|
||
print(f"\n📊 恢复结果: {done}✅ {running}🔄 (共{len(results)})")
|
||
if running > 0:
|
||
print(f"💡 还有 {running} 个镜头在生成中，稍后再次运行 batch-resume")
|
||
return results
|