196 lines
7.4 KiB
Python
196 lines
7.4 KiB
Python
"""
|
||
Seedance 视频生成客户端
|
||
使用 new-api 原生格式：POST /v1/video/generations → GET /v1/video/generations/{task_id}
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
from typing import Any, Callable, Dict, Optional
|
||
|
||
import aiohttp
|
||
|
||
from ..utils.config import get_api_key_or_raise
|
||
from ..utils.http_error import async_request_with_retry
|
||
from ..utils.video_task import (
|
||
check_interrupt,
|
||
extract_error_message,
|
||
extract_progress,
|
||
extract_status,
|
||
interruptible_sleep,
|
||
is_failure_status,
|
||
is_success_status,
|
||
run_with_interrupt,
|
||
)
|
||
|
||
|
||
class SeedanceClient:
|
||
"""Seedance 视频生成客户端（new-api 原生三段式）"""
|
||
|
||
# 提交任务
|
||
CREATE_ENDPOINT = "/v1/video/generations"
|
||
# 查询任务状态：{task_id} 占位
|
||
STATUS_ENDPOINT = "/v1/video/generations/{task_id}"
|
||
|
||
POLL_INITIAL_INTERVAL = 4 # 首次轮询等待秒数
|
||
POLL_MAX_INTERVAL = 15 # 最大轮询间隔秒数
|
||
|
||
# new-api 返回的成功状态值
|
||
SUCCESS_STATUSES = {"succeeded", "success", "completed", "done", "finished"}
|
||
FAILURE_STATUSES = {"failed", "fail", "error", "expired"}
|
||
|
||
def __init__(self):
|
||
self.api_key = get_api_key_or_raise()
|
||
self.base_url = "https://api.o1key.com"
|
||
|
||
def _headers(self) -> Dict[str, str]:
|
||
return {
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
|
||
# ── 1. 提交任务 ────────────────────────────────────────────────────
|
||
|
||
async def submit_async(
|
||
self,
|
||
body: Dict[str, Any],
|
||
session: aiohttp.ClientSession,
|
||
) -> str:
|
||
"""提交视频生成任务，返回 task_id"""
|
||
url = f"{self.base_url}{self.CREATE_ENDPOINT}"
|
||
check_interrupt()
|
||
resp = await run_with_interrupt(async_request_with_retry(
|
||
session, "POST", url, json=body, headers=self._headers(), prefix="Seedance 提交: "
|
||
))
|
||
check_interrupt()
|
||
text = await resp.text()
|
||
data = json.loads(text)
|
||
|
||
# new-api 返回字段：id / task_id
|
||
task_id = data.get("id") or data.get("task_id")
|
||
if not task_id:
|
||
raise RuntimeError(f"API 未返回任务 ID，响应：{data}")
|
||
return task_id
|
||
|
||
# ── 2. 轮询状态 ────────────────────────────────────────────────────
|
||
|
||
async def poll_async(
|
||
self,
|
||
task_id: str,
|
||
session: aiohttp.ClientSession,
|
||
on_progress: Optional[Callable[[int], None]] = None,
|
||
) -> str:
|
||
"""轮询任务状态，成功后返回视频 URL"""
|
||
url = f"{self.base_url}{self.STATUS_ENDPOINT.format(task_id=task_id)}"
|
||
interval = self.POLL_INITIAL_INTERVAL
|
||
|
||
while True:
|
||
check_interrupt()
|
||
async with session.get(url, headers=self._headers()) as resp:
|
||
text = await resp.text()
|
||
if resp.status != 200:
|
||
try:
|
||
err = json.loads(text)
|
||
msg = (err.get("error", {}).get("message")
|
||
or err.get("message")
|
||
or text)
|
||
except Exception:
|
||
msg = text
|
||
raise RuntimeError(f"状态查询失败 ({resp.status}): {msg}")
|
||
result = json.loads(text)
|
||
|
||
# new-api 包装格式：真实数据在 result["data"] 里
|
||
inner = result.get("data") or result
|
||
|
||
status = extract_status(result)
|
||
|
||
# 解析进度
|
||
progress_pct = extract_progress(result)
|
||
|
||
print(f"[Seedance] 生成中 {progress_pct}%")
|
||
if on_progress:
|
||
on_progress(progress_pct)
|
||
|
||
if is_success_status(status):
|
||
# 响应结构：result["data"] = inner，inner["data"] = platform_data
|
||
# 视频 URL 在 inner["result_url"] 或 inner["data"]["content"]["video_url"]
|
||
platform_data = inner.get("data") or {}
|
||
content = platform_data.get("content") or {}
|
||
video_url = (
|
||
inner.get("result_url")
|
||
or content.get("video_url")
|
||
or platform_data.get("video_url")
|
||
or inner.get("url")
|
||
)
|
||
if not video_url:
|
||
raise RuntimeError(f"任务成功但未找到视频 URL，响应：{result}")
|
||
# 末帧图片 URL 在 inner["data"]["content"]["last_frame_url"]
|
||
last_frame_url = (
|
||
content.get("last_frame_url")
|
||
or platform_data.get("last_frame_url")
|
||
or inner.get("last_frame_url")
|
||
)
|
||
return video_url, last_frame_url
|
||
|
||
if is_failure_status(status, result):
|
||
reason = extract_error_message(result)
|
||
raise RuntimeError(f"视频生成失败：{reason}")
|
||
|
||
await interruptible_sleep(interval)
|
||
interval = min(interval * 1.5, self.POLL_MAX_INTERVAL)
|
||
|
||
# ── 3. 下载视频 ────────────────────────────────────────────────────
|
||
|
||
async def download_async(
|
||
self,
|
||
video_url: str,
|
||
save_path: str,
|
||
session: aiohttp.ClientSession,
|
||
) -> str:
|
||
"""下载视频到本地，返回本地路径"""
|
||
print(f"[Seedance] 下载视频...")
|
||
check_interrupt()
|
||
async with session.get(video_url, allow_redirects=True) as resp:
|
||
if resp.status != 200:
|
||
raise RuntimeError(f"视频下载失败 ({resp.status})")
|
||
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
||
with open(save_path, "wb") as f:
|
||
async for chunk in resp.content.iter_chunked(8192):
|
||
check_interrupt()
|
||
f.write(chunk)
|
||
return save_path
|
||
|
||
# ── 全流程入口（供节点调用）────────────────────────────────────────
|
||
|
||
async def generate_async(
|
||
self,
|
||
body: Dict[str, Any],
|
||
save_path: str,
|
||
on_stage: Optional[Callable[[str], None]] = None,
|
||
on_progress: Optional[Callable[[int], None]] = None,
|
||
) -> tuple:
|
||
"""提交 → 轮询 → 下载，返回 (本地视频路径, 末帧图片URL或None)"""
|
||
connector = aiohttp.TCPConnector(ssl=False, force_close=True)
|
||
async with aiohttp.ClientSession(connector=connector) as session:
|
||
|
||
# 提交
|
||
check_interrupt()
|
||
if on_stage:
|
||
on_stage("submitting")
|
||
task_id = await self.submit_async(body, session)
|
||
print(f"[Seedance] 任务已提交 → {task_id}")
|
||
if on_stage:
|
||
on_stage(f"submitted:{task_id}")
|
||
|
||
# 轮询
|
||
video_url, last_frame_url = await self.poll_async(task_id, session, on_progress=on_progress)
|
||
|
||
# 下载
|
||
if on_stage:
|
||
on_stage("downloading")
|
||
path = await self.download_async(video_url, save_path, session)
|
||
|
||
if on_stage:
|
||
on_stage("done")
|
||
return path, last_frame_url
|