返回 VideoClaw
video_seedance.py
根目录 / video-claw / video-claw / backend / models / video_seedance.py
1 """
2 Seedance 视频生成 API 客户端 (字节跳动 ARK)
3
4 """
5
6 import os
7 import sys
8
9 models_dir = os.path.dirname(os.path.abspath(__file__))
10 backend_dir = os.path.dirname(models_dir)
11 if backend_dir not in sys.path:
12 sys.path.insert(0, backend_dir)
13
14 import time
15 import logging
16 import requests
17 import base64
18 from typing import Optional
19 from config import Config
20
21 logger = logging.getLogger(__name__)
22
23 class SeedanceVideoClient:
24 """
25 Seedance 视频生成客户端(字节跳动 ARK)
26 支持图生视频功能,采用 提交任务 -> 轮询 -> 下载 的异步流程
27 """
28
29 def __init__(
30 self,
31 api_key: Optional[str] = None,
32 base_url: Optional[str] = None,
33 timeout: int = 120,
34 ) -> None:
35 self.api_key = api_key or Config.ARK_API_KEY
36 self.base_url = (base_url or Config.ARK_BASE_URL or "https://ark.cn-beijing.volces.com/api/v3").rstrip("/")
37 self.timeout = timeout
38
39 if not self.api_key:
40 logger.warning("SeedanceVideoClient: ARK_API_KEY 未设置")
41
42 def _headers(self) -> dict:
43 return {
44 "Authorization": f"Bearer {self.api_key}",
45 "Content-Type": "application/json",
46 }
47
48 def generate_video(
49 self,
50 prompt: str,
51 image_path: str,
52 save_path: str,
53 model: str = "doubao-seedance-2-0-260128",
54 duration: int = 5,
55 **kwargs
56 ) -> str:
57 """
58 图生视频完整流程
59
60 Args:
61 prompt: 提示词
62 image_path: 输入图片本地路径
63 save_path: 输出视频保存路径
64 model: 模型名称
65 duration: 视频时长
66 """
67 if not self.api_key:
68 raise RuntimeError("ARK_API_KEY not set.")
69
70 # 1. 提交任务
71 task_id = self._submit_task(prompt, image_path, model, duration, **kwargs)
72
73 # 2. 轮询等待
74 video_url = self._poll_until_done(task_id)
75
76 # 3. 下载视频
77 self._download_video(video_url, save_path)
78
79 return video_url
80
81 def _submit_task(self, prompt: str, image_path: str, model: str, duration: int, **kwargs) -> str:
82 # 根据 Seedance 2.0 文档更新接口路径
83 url = f"{self.base_url}/contents/generations/tasks"
84
85 if not os.path.exists(image_path):
86 raise FileNotFoundError(f"输入图片不存在: {image_path}")
87
88 with open(image_path, "rb") as f:
89 img_data = base64.b64encode(f.read()).decode("utf-8")
90 ext = os.path.splitext(image_path)[1].lower()
91 mime = "image/png" if ext == ".png" else "image/jpeg"
92 image_base64 = f"data:{mime};base64,{img_data}"
93
94 # 构建 content 数组
95 content = []
96 if prompt:
97 content.append({
98 "type": "text",
99 "text": prompt
100 })
101
102 # 图生视频-首帧
103 content.append({
104 "type": "image_url",
105 "image_url": {
106 "url": image_base64
107 },
108 "role": "first_frame"
109 })
110
111 payload = {
112 "model": model,
113 "content": content,
114 "duration": duration,
115 "ratio": kwargs.get("ratio", "adaptive"),
116 "resolution": kwargs.get("resolution", "720p")
117 }
118
119 # 合并其他可选参数 (如 seed, watermark)
120 for key in ["seed", "watermark", "generate_audio"]:
121 if key in kwargs:
122 payload[key] = kwargs[key]
123
124 logger.info(f"SeedanceVideoClient: 提交任务 model={model}, duration={duration}s")
125 resp = requests.post(
126 url,
127 headers=self._headers(),
128 json=payload,
129 timeout=self.timeout,
130 proxies=Config.requests_proxies("ark"),
131 )
132
133 if not resp.ok:
134 logger.error(f"Seedance 提交失败: {resp.text}")
135 resp.raise_for_status()
136
137 data = resp.json()
138 task_id = data.get("id")
139 if not task_id:
140 raise RuntimeError(f"Seedance API 未返回任务 ID: {data}")
141
142 return task_id
143
144 def _poll_until_done(self, task_id: str, max_polls: int = 120, interval: int = 5) -> str:
145 # 同步更新查询接口路径
146 url = f"{self.base_url}/contents/generations/tasks/{task_id}"
147
148 for i in range(max_polls):
149 resp = requests.get(
150 url,
151 headers=self._headers(),
152 timeout=30,
153 proxies=Config.requests_proxies("ark"),
154 )
155 resp.raise_for_status()
156 data = resp.json()
157
158 status = data.get("status")
159 if status == "succeeded":
160 # 根据实际返回体,URL 位于 content.video_url 或 video_url
161 video_url = data.get("content", {}).get("video_url") or data.get("video_url")
162 if not video_url:
163 raise RuntimeError(f"Seedance 任务成功但未返回视频 URL: {data}")
164 return video_url
165 elif status in ("failed", "expired"):
166 error_msg = data.get("error", {}).get("message") or data.get("status_msg") or "未知错误"
167 raise RuntimeError(f"Seedance 视频生成{status}: {error_msg}")
168
169 logger.debug(f"SeedanceVideoClient: 任务进行中 {task_id}, status={status}, poll={i+1}")
170 time.sleep(interval)
171
172 raise TimeoutError(f"Seedance 视频生成超时 (task_id={task_id})")
173
174 def _download_video(self, url: str, save_path: str):
175 os.makedirs(os.path.dirname(save_path), exist_ok=True)
176 resp = requests.get(url, stream=True, timeout=120, proxies=Config.requests_proxies("ark"))
177 resp.raise_for_status()
178 with open(save_path, "wb") as f:
179 for chunk in resp.iter_content(chunk_size=8192):
180 if chunk:
181 f.write(chunk)
182 logger.info(f"SeedanceVideoClient: 视频已保存: {save_path}")
183
184 if __name__ == "__main__":
185 import sys
186 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
187 from config import Config
188
189 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
190
191 # ── 测试参数(按需修改) ──
192 # IMAGE_PATH = "code/result/image/test_avail/test_input.png"
193 IMAGE_PATH = "code/result/image/test_avail/test_input_human.jpg"
194 OUTPUT_PATH = "code/result/video/test_avail/seedance_test_output.mp4"
195 PROMPT = "女生把财务报表交给男生,男生看到后喜极而泣"
196 # MODELS = ["doubao-seedance-2-0-fast-260128", "doubao-seedance-2-0-260128"]
197 MODELS = ["doubao-seedance-2-0-fast-260128"]
198 DURATION = 5
199
200 print("=== Seedance (ARK) 图生视频测试 ===")
201 api_key = Config.ARK_API_KEY
202 base_url = Config.ARK_BASE_URL
203
204 if not api_key:
205 print("✗ ARK_API_KEY 未设置,请检查 config.yaml 配置")
206 sys.exit(1)
207
208 if not os.path.exists(IMAGE_PATH):
209 print(f"✗ 输入图片不存在: {IMAGE_PATH}")
210 sys.exit(1)
211
212 print(f" API Key : {api_key[:6]}***{api_key[-4:]}")
213 print(f" Base URL : {base_url}")
214
215 for model in MODELS:
216 print("\n" + "="*40)
217 print(f" 输入图片 : {IMAGE_PATH}")
218 print(f" 输出路径 : {OUTPUT_PATH}")
219 print(f" 模型 : {model}")
220 print(f" 时长 : {DURATION}s")
221 if PROMPT:
222 print(f" 提示词 : {PROMPT[:80]}")
223
224 try:
225 client = SeedanceVideoClient(api_key=api_key, base_url=base_url)
226 print("✓ 客户端初始化成功")
227
228 start = time.time()
229 video_url = client.generate_video(
230 prompt=PROMPT,
231 image_path=IMAGE_PATH,
232 save_path=OUTPUT_PATH,
233 model=model,
234 duration=DURATION,
235 )
236 elapsed = time.time() - start
237
238 print(f"✓ 视频生成完成!耗时 {elapsed:.1f}s")
239 print(f" 远端 URL : {video_url}")
240 print(f" 本地文件 : {os.path.abspath(OUTPUT_PATH)}")
241 print(f" 文件大小 : {os.path.getsize(OUTPUT_PATH) / 1024 / 1024:.2f} MB")
242 except Exception as e:
243 print(f"✗ 失败: {e}")
244 sys.exit(1)
245 break # 只测试第一个模型
246
246 lines PYTHON