MagicDriveDiT 使用教程

MagicDriveDiT 使用教程

1. 项目介绍

MagicDriveDiT 是一个基于 DiT 架构的开源项目，专注于为自动驾驶应用生成高质量的高分辨率长视频。它通过流匹配增强扩展性，并采用渐进式训练策略来处理复杂场景。该项目整合了空间时间条件编码，实现了对空间时间潜在变量的精确控制。实验证明，MagicDriveDiT 在生成真实街道场景视频方面具有优越的性能，帧数更多，分辨率更高，为自动驾驶的多种任务提供了扩展潜力。

2. 项目快速启动

环境搭建

首先，克隆项目仓库：

git clone https://github.com/flymin/MagicDriveDiT.git

根据你使用的服务器类型（NVIDIA 或 Ascend），设置相应的 Python 环境。以下是 NVIDIA 服务器的环境配置步骤：

确保你的环境中包含以下包：

• torch==2.4.0
• torchvision==0.19.0
• apex (从源代码构建或选择正确的轮包)
• xformers>=0.0.27
• flash-attn>=2.6.3

安装 ColossalAI：

git clone https://github.com/flymin/ColossalAI.git
git checkout pt2.4 && git pull
cd ColossalAI
BUILD_EXT=1 pip install .

安装其他依赖：

pip install -r requirements/requirements.txt

对于 Ascend 服务器，请参照相应的指南设置 PyTorch 环境，并安装必要的依赖。

准备预训练权重和数据

下载 VAE 和文本编码器的预训练权重，并按照以下结构组织：

${CODE_ROOT}/pretrained/
├── CogVideoX-2b
│   └── vae
└── t5-v1_1-xxl

下载 stage-3 检查点，并放置在 ${CODE_ROOT}/ckpts/ 目录下。

准备 nuScenes 数据集，并按照项目指南进行预处理。

运行推理

以下是在 GPU 服务器上运行模型推理的示例命令：

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
torchrun --standalone --nproc_per_node ${GPUS} scripts/inference_magicdrive.py ${CFG} \
--cfg-options model.from_pretrained=PATH_TO_MODEL

其中 ${GPUS} 是用于并行计算的 GPU 数量，${CFG} 是位于 configs/magicdrive/inference/ 目录下的配置文件，${PATH_TO_MODEL} 是检查点文件路径。

3. 应用案例和最佳实践

• 视频生成：使用 MagicDriveDiT 生成用于自动驾驶系统训练和测试的高分辨率长视频。
• 性能优化：通过调整配置文件中的参数，优化模型在特定硬件上的性能表现。

4. 典型生态项目

• 自动驾驶系统：集成 MagicDriveDiT 生成的视频数据，提高自动驾驶系统在不同环境和条件下的感知能力。
• 机器学习研究：研究人员可以利用 MagicDriveDiT 进行视频生成相关的研究和实验。