cang-ying/video-ai-system/pipeline.py

#!/usr/bin/env python3
"""
pipeline.py — 苍耳管线 · 统一生产入口
整合所有工具为一条命令链路：剧本→分镜→路由→出图→放大→拼接→配乐→成片

用法:
  python pipeline.py storyboard <剧本> -e 1          # ① 剧本→分镜
  python pipeline.py route <分镜.json>                 # ② 路由分配
  python pipeline.py render <分镜.json>                # ③ 批量出图
  python pipeline.py upscale <图片文件夹>               # ④ 高清放大
  python pipeline.py compose <文件夹> -o EP01.mp4     # ⑤ 剪辑拼接
  python pipeline.py assets <项目目录>                 # ⑥ 资产生成
  
  # 一键全自动（需要API密钥）
  python pipeline.py go <剧本> -e 1                    # ⑦ 全自动一条龙
"""
import sys, os, json, argparse, subprocess

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

def cmd_storyboard(args):
    """① 剧本→分镜"""
    sys.argv = ['run_storyboard.py', args.script, '-e', str(args.episode)]
    if args.pro:
        sys.argv.append('--pro')
    if args.dry_run:
        sys.argv.append('--dry-run')
    if args.output:
        sys.argv.extend(['-o', args.output])
    
    from tools.run_storyboard import main
    main()

def cmd_route(args):
    """② 路由分配"""
    sys.argv = ['route_shots.py', args.storyboard]
    if args.dry_run:
        sys.argv.append('--dry-run')
    if args.summary:
        sys.argv.append('--summary')
    if args.output:
        sys.argv.extend(['-o', args.output])
    
    from tools.route_shots import main
    main()

def cmd_render(args):
    """③ 分镜→出图桥接"""
    sys.argv = ['storyboard_to_workflow.py', args.storyboard]
    if args.output_dir:
        sys.argv.extend(['-o', args.output_dir])
    if args.style:
        sys.argv.extend(['--style', args.style])
    if args.generate_sh:
        sys.argv.append('--generate-sh')
    if args.workflow:
        sys.argv.extend(['--workflow', args.workflow])
    
    from tools.storyboard_to_workflow import main
    main()

def cmd_upscale(args):
    """④ 高清放大"""
    from tools.storyboard_to_workflow import main as sb_main
    print('🔍 高清放大: 可使用 ComfyUI 蓝图 Image Upscale(Z-image-Turbo).json')
    print(f'   打开 ComfyUI → 加载蓝图 → Image Upscale')
    print(f'   输入图片: {args.input_dir}')
    print()
    print('   或运行已生成的工作流JSON (如果有):')
    print(f'   ls {args.input_dir}/workflows/')

def cmd_compose(args):
    """⑤ 剪辑拼接"""
    sys.argv = ['video_composer.py', args.folder, '-o', args.output]
    if args.title:
        sys.argv.extend(['--title', args.title])
    if args.credits:
        sys.argv.extend(['--credits', args.credits])
    if args.subtitle:
        sys.argv.extend(['--subtitle', args.subtitle])
    if args.crossfade:
        sys.argv.extend(['--crossfade', str(args.crossfade)])
    
    from tools.video_composer import main
    main()

def cmd_assets(args):
    """⑥ 资产生成"""
    sys.argv = ['generate_assets_local.py', args.project_dir]
    if args.list:
        sys.argv.append('--list')
    if args.dry_run:
        sys.argv.append('--dry-run')
    if args.comfy_url:
        sys.argv.extend(['--comfy-url', args.comfy_url])
    
    from tools.generate_assets_local import main
    main()

def cmd_go(args):
    """⑦ 一键全自动（剧本→成片）"""
    script_path = args.script
    episode = args.episode
    
    print(f'🚀 苍耳管线 · 全自动模式')
    print(f'📖 剧本: {script_path}')
    print(f'🎬 集号: EP{episode:02d}')
    print()
    
    # 检查文件
    if not os.path.isfile(script_path):
        print(f'❌ 找不到剧本: {script_path}')
        sys.exit(1)
    
    # Step 1: 剧本→分镜
    print('═' * 40)
    print('① 剧本→分镜')
    print('═' * 40)
    sys.argv = ['run_storyboard.py', script_path, '-e', str(episode)]
    if args.pro:
        sys.argv.append('--pro')
    
    from tools.run_storyboard import main as sb_main
    try:
        sb_main()
    except SystemExit:
        pass
    
    # 找生成的分镜JSON
    script_dir = os.path.dirname(os.path.abspath(script_path))
    script_base = os.path.splitext(os.path.basename(script_path))[0]
    sb_json = None
    for f in os.listdir(script_dir):
        if f.startswith('STORYBOARD') and f.endswith('.json'):
            sb_json = os.path.join(script_dir, f)
    
    if not sb_json:
        print('❌ 找不到生成的分镜JSON')
        sys.exit(1)
    
    print(f'\n✅ 分镜: {sb_json}')
    
    # Step 2: 路由分配
    print()
    print('═' * 40)
    print('② 逐镜路由分配')
    print('═' * 40)
    sys.argv = ['route_shots.py', sb_json]
    from tools.route_shots import main as rt_main
    try:
        rt_main()
    except SystemExit:
        pass
    
    # Step 3: 生成出图工作流
    print()
    print('═' * 40)
    print('③ 生成出图工作流')
    print('═' * 40)
    sys.argv = ['storyboard_to_workflow.py', sb_json]
    if args.style:
        sys.argv.extend(['--style', args.style])
    sys.argv.append('--generate-sh')
    from tools.storyboard_to_workflow import main as sw_main
    try:
        sw_main()
    except SystemExit:
        pass
    
    print()
    print('═' * 40)
    print('🎯 完成! 后续步骤:')
    print('═' * 40)
    print(f'   📄 分镜: {sb_json}')
    print(f'   📋 路由: 已写入分镜JSON')
    print(f'   🎨 出图: 启动ComfyUI后运行渲染脚本')
    print(f'   ✂️ 剪辑: python pipeline.py compose <输出目录> -o EP{episode:02d}.mp4')

def main():
    parser = argparse.ArgumentParser(description='苍耳管线 · 统一生产入口')
    parser.add_argument('command', nargs='?', default='help',
                       help='命令: storyboard|route|render|upscale|compose|assets|go')
    parser.add_argument('args', nargs='*', help='命令参数')
    
    # 先解析command
    known_commands = ['storyboard', 'route', 'render', 'upscale', 'compose', 'assets', 'go', 'help']
    
    if len(sys.argv) < 2 or sys.argv[1] not in known_commands:
        print(__doc__)
        return
    
    cmd = sys.argv[1]
    
    if cmd == 'help':
        print(__doc__)
        return
    
    # 移除第一个参数(command)，让子命令的argparse自己解析
    # 但我们需要根据command来设置子解析器
    if cmd == 'storyboard':
        p = argparse.ArgumentParser()
        p.add_argument('script', help='剧本文件路径')
        p.add_argument('-e', '--episode', type=int, default=1, help='集号')
        p.add_argument('--pro', action='store_true', help='使用Pro模型')
        p.add_argument('-o', '--output', help='输出路径')
        p.add_argument('--dry-run', action='store_true', help='模拟运行')
        args = p.parse_args(sys.argv[2:])
        cmd_storyboard(args)
    
    elif cmd == 'route':
        p = argparse.ArgumentParser()
        p.add_argument('storyboard', help='分镜JSON路径')
        p.add_argument('-o', '--output', help='输出路径')
        p.add_argument('--dry-run', action='store_true', help='模拟运行')
        p.add_argument('--summary', action='store_true', help='只统计')
        args = p.parse_args(sys.argv[2:])
        cmd_route(args)
    
    elif cmd == 'render':
        p = argparse.ArgumentParser()
        p.add_argument('storyboard', help='分镜JSON路径')
        p.add_argument('-o', '--output-dir', help='输出目录')
        p.add_argument('--style', default='', help='全局风格')
        p.add_argument('--generate-sh', action='store_true', help='生成shell脚本')
        p.add_argument('--workflow', help='基础工作流模板')
        args = p.parse_args(sys.argv[2:])
        cmd_render(args)
    
    elif cmd == 'upscale':
        p = argparse.ArgumentParser()
        p.add_argument('input_dir', help='图片文件夹')
        args = p.parse_args(sys.argv[2:])
        cmd_upscale(args)
    
    elif cmd == 'compose':
        p = argparse.ArgumentParser()
        p.add_argument('folder', help='视频片段文件夹')
        p.add_argument('-o', '--output', default='output.mp4', help='输出文件')
        p.add_argument('--title', help='片头文字')
        p.add_argument('--credits', help='片尾文字')
        p.add_argument('--subtitle', help='字幕文件')
        p.add_argument('--crossfade', type=int, default=0, help='交叉淡入淡出')
        args = p.parse_args(sys.argv[2:])
        cmd_compose(args)
    
    elif cmd == 'assets':
        p = argparse.ArgumentParser()
        p.add_argument('project_dir', help='项目目录')
        p.add_argument('--comfy-url', help='ComfyUI API地址')
        p.add_argument('--list', action='store_true', help='只列出')
        p.add_argument('--dry-run', action='store_true', help='模拟运行')
        args = p.parse_args(sys.argv[2:])
        cmd_assets(args)
    
    elif cmd == 'go':
        p = argparse.ArgumentParser()
        p.add_argument('script', help='剧本文件路径')
        p.add_argument('-e', '--episode', type=int, default=1, help='集号')
        p.add_argument('--pro', action='store_true', help='使用Pro模型')
        p.add_argument('--style', default='', help='全局风格')
        args = p.parse_args(sys.argv[2:])
        cmd_go(args)

if __name__ == '__main__':
    main()