邓文兵
2c96bf7b7f
- 添加视频下载重试机制，最多重试3次 - 增加下载超时设置为60秒 - 添加下载失败时的错误提示信息 - 修复任务状态检查逻辑，避免空响应导致的错误 - 改进错误处理流程，提高脚本稳定性
157 lines
5.6 KiB
Python
157 lines
5.6 KiB
Python
"""即梦 AI 视频生成工具
|
||
|
||
支持：
|
||
- 文生视频 (t2v): jimeng_t2v_v30 (720P) / jimeng_t2v_v30_1080p (1080P)
|
||
- 图生视频-首帧 (i2v): jimeng_i2v_first_v30 (720P)
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import sys
|
||
import time
|
||
import base64
|
||
import requests
|
||
|
||
from jimeng_gen import jimeng_request, _ensure_dir
|
||
|
||
|
||
# req_key 映射
|
||
REQ_KEYS = {
|
||
't2v_720p': 'jimeng_t2v_v30',
|
||
't2v_1080p': 'jimeng_t2v_v30_1080p',
|
||
'i2v_720p': 'jimeng_i2v_first_v30',
|
||
}
|
||
|
||
|
||
def generate_video_t2v(prompt, output_path, resolution='720p', aspect_ratio='16:9', frames=121, seed=-1):
|
||
"""文生视频：提交任务并轮询获取结果"""
|
||
req_key = REQ_KEYS.get(f't2v_{resolution}', REQ_KEYS['t2v_720p'])
|
||
print(f"[即梦视频] 文生视频 ({resolution}): {prompt[:50]}...")
|
||
|
||
body = {
|
||
'req_key': req_key,
|
||
'prompt': prompt,
|
||
'frames': frames,
|
||
'aspect_ratio': aspect_ratio,
|
||
'seed': seed,
|
||
}
|
||
|
||
return _submit_and_poll(req_key, body, output_path)
|
||
|
||
|
||
def generate_video_i2v(prompt, image_path_or_url, output_path, aspect_ratio='16:9', frames=121, seed=-1):
|
||
"""图生视频-首帧：输入首帧图片和提示词生成视频"""
|
||
req_key = REQ_KEYS['i2v_720p']
|
||
print(f"[即梦视频] 图生视频-首帧: {prompt[:50]}...")
|
||
|
||
body = {
|
||
'req_key': req_key,
|
||
'prompt': prompt,
|
||
'frames': frames,
|
||
'seed': seed,
|
||
}
|
||
|
||
# 判断是 URL 还是本地文件
|
||
if image_path_or_url.startswith('http'):
|
||
body['image_urls'] = [image_path_or_url]
|
||
else:
|
||
with open(image_path_or_url, 'rb') as f:
|
||
img_b64 = base64.b64encode(f.read()).decode('utf-8')
|
||
body['binary_data_base64'] = [img_b64]
|
||
|
||
return _submit_and_poll(req_key, body, output_path)
|
||
|
||
|
||
def _submit_and_poll(req_key, body, output_path, max_wait=180, poll_interval=10):
|
||
"""提交任务并轮询结果，下载视频到本地"""
|
||
# 1. 提交任务
|
||
submit_resp = jimeng_request('CVSync2AsyncSubmitTask', body)
|
||
|
||
if submit_resp.get('code') != 10000:
|
||
print(f"[即梦视频] 提交失败: {submit_resp}")
|
||
return None
|
||
|
||
task_id = submit_resp['data']['task_id']
|
||
print(f"[即梦视频] 任务已提交, task_id: {task_id}")
|
||
|
||
# 2. 轮询查询结果
|
||
max_polls = max_wait // poll_interval
|
||
for i in range(max_polls):
|
||
time.sleep(poll_interval)
|
||
result = jimeng_request('CVSync2AsyncGetResult', {
|
||
'req_key': req_key,
|
||
'task_id': task_id,
|
||
})
|
||
|
||
resp_code = result.get('code')
|
||
resp_data = result.get('data', {})
|
||
|
||
if resp_code == 10000 and resp_data.get('status') == 'done':
|
||
video_url = resp_data.get('video_url')
|
||
if video_url:
|
||
# 下载视频（带重试）
|
||
print(f"[即梦视频] 生成完成，正在下载...")
|
||
video_data = None
|
||
for attempt in range(3):
|
||
try:
|
||
video_data = requests.get(video_url, timeout=60, verify=False).content
|
||
break
|
||
except Exception as e:
|
||
print(f"[即梦视频] 下载重试 ({attempt+1}/3): {e}")
|
||
import time as _t; _t.sleep(3)
|
||
if video_data is None:
|
||
print(f"[即梦视频] 下载失败，视频URL: {video_url[:200]}")
|
||
return None
|
||
_ensure_dir(output_path)
|
||
with open(output_path, 'wb') as f:
|
||
f.write(video_data)
|
||
print(f"[即梦视频] 视频已保存: {output_path}")
|
||
return output_path
|
||
else:
|
||
print(f"[即梦视频] 任务完成但未找到视频URL: {json.dumps(result, ensure_ascii=False)[:300]}")
|
||
return None
|
||
|
||
if resp_code != 10000 and (resp_data is None or resp_data.get('status') not in (None, 'running', 'pending')):
|
||
print(f"[即梦视频] 任务失败: {json.dumps(result, ensure_ascii=False)[:300]}")
|
||
return None
|
||
|
||
print(f"[即梦视频] 等待中... ({i+1}/{max_polls})")
|
||
|
||
print("[即梦视频] 超时")
|
||
return None
|
||
|
||
|
||
def main():
|
||
import argparse
|
||
parser = argparse.ArgumentParser(description='即梦 AI 视频生成')
|
||
parser.add_argument('prompt', help='视频描述提示词')
|
||
parser.add_argument('output', help='输出视频路径 (.mp4)')
|
||
parser.add_argument('--mode', choices=['t2v', 'i2v'], default='t2v', help='生成模式: t2v=文生视频, i2v=图生视频')
|
||
parser.add_argument('--resolution', choices=['720p', '1080p'], default='720p', help='分辨率 (仅 t2v)')
|
||
parser.add_argument('--aspect-ratio', default='16:9', help='宽高比: 16:9, 4:3, 1:1, 3:4, 9:16, 21:9')
|
||
parser.add_argument('--frames', type=int, default=121, help='视频帧数 (默认121)')
|
||
parser.add_argument('--image', help='首帧图片路径或URL (仅 i2v 模式)')
|
||
parser.add_argument('--seed', type=int, default=-1, help='随机种子 (-1=随机)')
|
||
|
||
args = parser.parse_args()
|
||
|
||
if args.mode == 'i2v':
|
||
if not args.image:
|
||
parser.error('i2v 模式需要 --image 参数')
|
||
result = generate_video_i2v(args.prompt, args.image, args.output,
|
||
aspect_ratio=args.aspect_ratio, frames=args.frames, seed=args.seed)
|
||
else:
|
||
result = generate_video_t2v(args.prompt, args.output,
|
||
resolution=args.resolution, aspect_ratio=args.aspect_ratio,
|
||
frames=args.frames, seed=args.seed)
|
||
|
||
if result:
|
||
print(f"\n视频生成成功: {result}")
|
||
else:
|
||
print("\n视频生成失败")
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|