282 lines
12 KiB
Python
Raw Normal View History

# 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