cang-ying/video-ai-system/tools/run_storyboard.py

118 lines
4.3 KiB
Python
Raw Normal View History

#!/usr/bin/env python3
# 步骤⑦ storyboard 引擎: 豆包 AI 拆解剧本 → 结构化分镜 JSON
# 通用版 — 支持任意剧本文件
# DEPENDS: lib/doubao_chat.py
import sys, os, json, argparse
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from lib.doubao_chat import breakdown_script
def extract_json(text):
"""从 AI 回复中提取 JSON 内容(兼容带 ```json 包裹的情况)"""
if '```json' in text:
return text.split('```json')[1].split('```')[0].strip()
if '```' in text:
return text.split('```')[1].split('```')[0].strip()
# 尝试直接解析
maybe = text.strip()
if maybe.startswith('{'):
return maybe
return maybe
def main():
parser = argparse.ArgumentParser(description='豆包AI → 拆解剧本为结构化分镜JSON')
parser.add_argument('script', help='剧本文件路径 (.hdlp 或 .txt)')
parser.add_argument('-e', '--episode', type=int, default=1, help='集号 (默认: 1)')
parser.add_argument('-o', '--output', help='输出路径 (默认: 自动推导)')
parser.add_argument('--pro', action='store_true', help='使用 Pro 模型 (doubao-pro, 更强更贵)')
parser.add_argument('--dry-run', action='store_true', help='只打印将发送的提示词不调API')
args = parser.parse_args()
# 读剧本
if not os.path.isfile(args.script):
print(f'❌ 找不到剧本文件: {args.script}')
sys.exit(1)
with open(args.script, 'r', encoding='utf-8') as f:
script_text = f.read()
print(f'📖 剧本: {args.script}')
print(f'📐 长度: {len(script_text)} 字符')
print(f'🎬 集号: EP{args.episode:02d}')
print(f'🧠 模型: {"doubao-pro" if args.pro else "doubao-lite"}')
print()
if args.dry_run:
print('🧪 DRY RUN — 将发送的提示词:')
print('=' * 60)
from lib.doubao_chat import breakdown_script as bd
print(bd.__doc__)
print('=' * 60)
return
# 调用豆包 API
print('🚀 发送豆包 AI 拆解剧本...')
result = breakdown_script(script_text, args.episode, model='pro' if args.pro else 'lite')
if 'error' in result:
print('❌ API ERROR:', json.dumps(result, ensure_ascii=False, indent=2))
sys.exit(1)
if 'content' not in result or not result['content']:
print('❌ API 返回空内容:', json.dumps(result, ensure_ascii=False, indent=2))
sys.exit(1)
content = result['content']
json_str = extract_json(content)
# 解析 JSON 验证有效性
try:
parsed = json.loads(json_str)
shot_count = len(parsed.get('shots', []))
except json.JSONDecodeError as e:
print(f'⚠️ JSON 解析失败: {e}')
print('原始内容:')
print(content[:500])
shot_count = 0
parsed = None
# 确定输出路径
if args.output:
out_path = args.output
else:
# 从剧本路径自动推导
script_dir = os.path.dirname(os.path.abspath(args.script))
script_base = os.path.splitext(os.path.basename(args.script))[0]
# 如果剧本文件名包含 SCRIPT 或 BREAKDOWN替换为 STORYBOARD
if 'SCRIPT' in script_base.upper() or 'BREAKDOWN' in script_base.upper():
storyboard_name = script_base.replace('SCRIPT', 'STORYBOARD').replace('BREAKDOWN', 'STORYBOARD')
else:
model_tag = 'PRO' if args.pro else 'AI'
storyboard_name = f'STORYBOARD-{model_tag}-EP{args.episode:02d}'
out_path = os.path.join(script_dir, f'{storyboard_name}.json')
# 保存
with open(out_path, 'w', encoding='utf-8') as f:
f.write(json_str if parsed else content.strip())
print(f'\n✅ 分镜已保存: {out_path}')
if parsed:
print(f'📊 共 {shot_count}')
print(f'💰 tokens: {result.get("tokens", "N/A")}')
print(f'⏱️ 完成原因: {result.get("finish", "N/A")}')
# 打印前3镜预览
if parsed and shot_count > 0:
print('\n📋 前3镜预览:')
for s in parsed['shots'][:3]:
print(f' {s.get("shot_number","?"):>4} | {s.get("camera","?"): <4} | {s.get("duration","?"):>2}s | {s.get("description","")[:40]}')
if __name__ == '__main__':
main()