邓文兵
9408daeeff
- 新增 scripts/jimeng_video_gen.py 脚本支持文生视频和图生视频 - 实现 t2v 和 i2v 两种生成模式，支持多种分辨率和宽高比 - 添加轮询任务状态和视频下载功能 - 更新即梦AI文档，补充视频生成API接口说明 - 新增 .claude/commands/generate-video.md 命令使用指南 - 扩展即梦API文档，包含视频生成参数配置和调用示例
147 lines
5.1 KiB
Python
147 lines
5.1 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 = requests.get(video_url).content
|
|
_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.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()
|