Video-BLADE: Block-Sparse Attention Meets Step Distillation for Efficient Video Generation

📖 Paper | 🚀 Homepage | 💾 Models

Video-BLADE是一个无需视频数据训练的高效视频生成框架，通过联合训练自适应稀疏注意力和步数蒸馏技术，实现了视频生成模型的显著加速。该项目实现了块稀疏注意力机制与步数蒸馏技术的结合，在保持生成质量的同时将推理步数从50步降低到8步。

📢 News

[2025-08] 🎉 Video-BLADE代码和预训练模型发布！
[2025-08] 📝 支持CogVideoX-5B和WanX-1.3B两种主流视频生成模型
[2025-08] ⚡ 实现8步高质量视频生成，相比50步baseline显著提速

✨ 主要特性

🚀 高效推理: 将推理步数从50步减少到8步，保持生成质量
🎯 自适应稀疏注意力: 块稀疏注意力机制，显著降低计算复杂度
📈 步数蒸馏: TDM(Trajectory Distillation Method)技术，无需视频数据即可训练
🎮 即插即用: 支持CogVideoX-5B和WanX-1.3B模型，无需修改原始架构

🛠️ 环境配置

系统要求

Python >= 3.11 (建议)
CUDA >= 11.6 (建议)
GPU内存 >= 24GB (推理)
GPU内存 >= 80GB (训练)

安装步骤

克隆仓库

git clone https://github.com/Tacossp/VIDEO-BLADE
cd VIDEO-BLADE

安装依赖包

# 使用uv安装依赖 (推荐)
uv pip install -r requirements.txt

# 或使用pip
pip install -r requirements.txt

编译Block-Sparse-Attention库

git clone https://github.com/mit-han-lab/Block-Sparse-Attention.git
cd Block-Sparse-Attention
pip install packaging
pip install ninja
python setup.py install
cd ..

📥 模型权重下载

基础模型权重

请下载以下基础模型权重并放置在指定目录：

CogVideoX-5B模型

# 从Hugging Face下载
git lfs install
git clone https://huggingface.co/zai-org/CogVideoX-5b cogvideox/CogVideoX-5b

WanX-1.3B模型

# 从Hugging Face下载
git clone https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers wanx/wan1.3b

预训练的Video-BLADE权重

我们提供了已经训练好的Video-BLADE权重：

# 下载预训练权重
git clone https://huggingface.co/GYP666/VIDEO-BLADE pretrained_weights

权重目录结构

确保您的权重目录结构如下：

VIDEO-BLADE/
├── cogvideox/
│   └── CogVideoX-5b/           # CogVideoX基础模型权重
├── wanx/
│   └── wan1.3b/               # WanX基础模型权重
└── pretrained_weights/         # Video-BLADE预训练权重
    ├── BLADE_cogvideox_weight/
    └── BLADe_wanx_weight/

🚀 快速开始 - 推理使用

CogVideoX推理

cd cogvideox
python train/inference.py \
    --lora_path ../pretrained_weights/ 
    cogvideox_checkpoints/your_checkpoint \
    --gpu 0

参数说明:

--lora_path: LoRA权重文件路径
--gpu: 使用的GPU设备ID (默认: 0)

输出: 生成的视频将保存在 cogvideox/outputs/inference/ 目录

WanX推理

cd wanx
python train/inference.py \
    --lora_path ../pretrained_weights/wanx_checkpoints/your_checkpoint \
    --gpu 0

输出: 生成的视频将保存在 wanx/outputs/ 目录

🔧 训练流程

第一步: Prompts预处理

在训练前，需要预处理提示词生成embeddings：

CogVideoX预处理

cd utils
python process_prompts_cogvideox.py \
    --input_file your_prompts.txt \
    --output_dir ../cogvideox/prompts \
    --model_path ../cogvideox/CogVideoX-5b \
    --batch_size 32 \
    --save_separate

参数说明:

--input_file: 包含prompts的txt文件，每行一个prompt
--output_dir: 输出embeddings的目录
--model_path: CogVideoX模型路径
--batch_size: 处理批次大小
--save_separate: 是否将每个embedding单独保存

WanX预处理

cd utils
python process_prompts_wanx.py

此脚本会自动处理 utils/all_dimension_aug_wanx.txt 中的prompts并生成相应的embeddings。

第二步: 启动训练

CogVideoX训练

cd cogvideox
bash train_tdm_1.sh

核心训练参数:

#如果不是8卡训练需要修改CUDA_VISIBLE_DEVICES和config.yaml的num_processes
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch \
    --config_file train/config.yaml \
    train/train_cogvideo_tdm.py \
    --pretrained_model_name_or_path CogVideoX-5b \        # 基础模型路径
    --mixed_precision bf16 \                              # 混合精度训练，减少显存使用
    --train_batch_size 5 \                                # 训练批次大小
    --gradient_accumulation_steps 4 \                     # 梯度累积步数
    --learning_rate 1e-4 \                                # student学习率
    --learning_rate_g 1e-4 \                              
    --learning_rate_fake 5e-4 \                           # fake model学习率
    --lambda_reg 0.5 \                                    # 正则化权重
    --k_step 8 \                                          # 蒸馏目标步数
    --cfg 3.5 \                                           # CFG引导强度
    --eta 0.9 \                                           # ETA参数
    --use_sparsity true \                                 # 启用稀疏注意力
    --rank 64 \
    --lora_alpha 64 \                                     # LoRA配置
    --max_train_steps 300 \                               # 最大训练步数
    --checkpointing_steps 15 \                            # 检查点保存间隔
    --gradient_checkpointing \                            # 梯度检查点，节省显存
    --enable_slicing \
    --enable_tiling   \                                   # VAE内存优化

WanX训练

cd wanx
bash train_wanx_tdm.sh

📊 项目结构

VIDEO-BLADE/
├── README.md                   # 项目说明文档
├── requirements.txt           # Python依赖列表
│
├── cogvideox/                 # CogVideoX相关代码
│   ├── CogVideoX-5b/         # 基础模型权重目录
│   ├── train/                # 训练脚本
│   │   ├── inference.py      # 推理脚本
│   │   ├── train_cogvideo_tdm.py  # 训练脚本
│   │   ├── train_tdm_1.sh    # 训练启动脚本
│   │   ├── modify_cogvideo.py # 模型修改脚本
│   │   └── config.yaml       # 训练配置文件
│   ├── prompts/              # 预处理的prompts和embeddings
│   └── outputs/              # 训练和推理输出
│
├── wanx/                     # WanX相关代码  
│   ├── wan1.3b/             # 基础模型权重目录
│   ├── train/               # 训练脚本
│   │   ├── inference.py     # 推理脚本
│   │   ├── train_wanx_tdm.py # 训练脚本
│   │   ├── train_wanx_tdm.sh # 训练启动脚本
│   │   └── modify_wan.py    # 模型修改脚本
│   ├── prompts/             # 预处理的prompts和embeddings
│   └── outputs/             # 训练和推理输出
│
├── utils/                   # 工具脚本
│   ├── process_prompts_cogvideox.py  # CogVideoX数据预处理
│   ├── process_prompts_wanx.py       # WanX数据预处理
│   └── all_dimension_aug_wanx.txt    # WanX训练prompts
│
├── Block-Sparse-Attention/  # 稀疏注意力库
│   ├── setup.py            # 编译安装脚本
│   ├── block_sparse_attn/  # 核心库代码
│   └── README.md           # 库使用说明
│
└── ds_config.json          # DeepSpeed配置文件

🤝 致谢

FlashAttention,Block-Sparse-Attention: 稀疏注意力实现基础
CogVideoX,Wan2.1: 模型支持
TDM:蒸馏实现基础
Diffusers: 扩散模型工具库

📄 引用

如果您在研究中使用了Video-BLADE，请引用我们的工作：

@misc{gu2025videobladeblocksparseattentionmeets,
    title={Video-BLADE: Block-Sparse Attention Meets Step Distillation for Efficient Video Generation},
    author={Youping Gu and Xiaolong Li and Yuhao Hu and Bohan Zhuang},
    year={2025},
    eprint={2508.10774},
    archivePrefix={arXiv},
    primaryClass={cs.CV},
    url={https://arxiv.org/abs/2508.10774},
}

📧 联系方式

如有问题或建议，欢迎：

Please contact Youping Gu ([email protected]) if you have any questions about this work.
提交issue: Github issue