MagicDrive-V2

本仓库包含论文的实现：

MagicDrive-V2：基于自适应控制的自动驾驶高分辨率长视频生成
高瑞源（https://gaoruiyuan.com/）¹, 陈凯（https://kaichen1998.github.io/）², 肖博（https://www.linkedin.com/in/bo-xiao-19909955/?originalSubdomain=ie）³, 洪兰青（https://scholar.google.com.sg/citations?user=2p7x6OUAAAAJ&hl=en）⁴, 李振国（https://scholar.google.com/citations?user=XboZC1AAAAAJ&hl=en）⁴, 徐强（https://cure-lab.github.io/）^1
1香港中文大学 ²香港科技大学 ³华为云 ⁴华为诺亚方舟实验室

https://github.com/user-attachments/assets/4c915eca-6583-422e-9a1a-1eebb9379ff4

摘要

简而言之 MagicDrive-V2 能够生成具有多样化 3D 几何控制和多视角一致性的高分辨率、长时长街景视频。

扩散模型的快速发展极大地推动了视频合成技术的进步，尤其是在可控视频生成领域，这对于自动驾驶等应用至关重要。然而，现有方法在可扩展性以及控制条件的整合方式上存在局限性，难以满足自动驾驶场景中对高分辨率、长时长视频的需求。本文我们提出了基于 DiT 架构的全新方法——MagicDrive-V2，以应对这些挑战。我们的方法通过流匹配提升了可扩展性，并采用渐进式训练策略来处理复杂场景。通过引入时空条件编码，MagicDrive-V2 实现了对时空潜在变量的精确控制。全面的实验表明，该方法在生成更高分辨率、更多帧数的真实街景视频方面表现出色。MagicDrive-V2 显著提升了视频生成质量和时空控制能力，从而拓展了其在自动驾驶相关任务中的应用潜力。

最新消息

• [2025年6月26日] MagicDrive-V2 被 ICCV 2025 接收 :tada: 。为与论文名称保持一致，我们将此仓库更名为 MagicDrive-V2。
• [2025年1月27日] 我们在项目页面更新了 Waymo 开放数据集上的微调结果。立即查看!
• [2024年12月7日] 第三阶段检查点及用于训练和推理的 nuScenes 元数据发布！
• [2024年12月3日] 训练与推理代码发布！我们稍后会更新 README 中的链接。
• [2024年11月22日] 论文和项目页面正式发布！请访问 https://gaoruiyuan.com/MagicDrive-V2/

待办事项

• 训练与推理代码
• 第三阶段预训练权重及 nuScenes 元数据
• 第一、二阶段预训练权重（后续发布）
• 使用 Waymo 数据集进行升级

快速入门

环境搭建

克隆本仓库

git clone https://github.com/flymin/MagicDrive-V2.git

代码已在 A800/H20/Ascend 910b 服务器上测试通过。请按照以下步骤设置 Python 环境：

[!注意]
请使用 pip 来配置环境。我们不建议直接使用 conda + yaml 进行环境配置。

NVIDIA 服务器 逐步指南：

1. 确保您的环境中已安装以下包：

   torch==2.4.0
   torchvision==0.19.0
   
   # 可能需要从源码编译
   apex (https://github.com/NVIDIA/apex)
   
   # 选择合适的轮子包或从源码编译
   xformers>=0.0.27
   flash-attn>=2.6.3
   

2. 安装 ColossalAI

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

3. 安装其他依赖

   pip install -r requirements/requirements.txt
   

更多详细信息请参考以下 YAML 文件：

• A800：requirements/a800_cu118.yaml
• H20：requirements/h20_cu124.yaml

Ascend 服务器 逐步指南：

1. 确保您的环境中已安装以下包（请参考 此页面 设置 PyTorch 环境）：

   # 基于 CANN 8.0RC2
   torch==2.3.1
   torchvision==0.18.1
   torch-npu==2.3.1
   apex (https://gitee.com/ascend/apex)
   
   # 选择合适的轮子包或从源码编译
   xformers==0.0.27
   

2. 安装 ColossalAI

   # 我们移除了对 `bitsandbytes` 的依赖。
   git clone https://github.com/flymin/ColossalAI.git
   git checkout ascend && git pull
   cd ColossalAI
   BUILD_EXT=1 pip install .
   

3. 安装其他依赖

   pip install -r requirements/requirements.txt
   

更多详细信息请参阅 requirements/910b_cann8.0.RC2_aarch64.yaml。

预训练权重

VAE：我们使用来自 THUDM/CogVideoX-2b 的 3DVAE。您只需下载 vae 子文件夹即可。

文本编码器：我们使用来自 google/t5-v1_1-xxl 的 T5 编码器。

请按以下结构组织这些文件：

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

MagicDrive-V2 检查点

请从 flymin/MagicDriveDiT-stage3-40k-ft 下载第三阶段检查点，并将其放置在 ${CODE_ROOT}/ckpts/ 目录下，如下所示：

${CODE_ROOT}/ckpts/
└── MagicDriveDiT-stage3-40k-ft

准备数据

我们按照 bevfusion 的说明 类似的方式准备 nuScenes 数据集。具体步骤如下：

1. 从 nuScenes 官网 下载 nuScenes 数据集，并将其放置在 ./data/ 目录下。你应该会看到以下文件结构：

   ${CODE_ROOT}/data/nuscenes
   ├── can_bus
   ├── maps
   ├── mini
   ├── samples
   ├── sweeps
   ├── v1.0-mini
   └── v1.0-trainval
   

2. 从 flymin/MagicDriveDiT-nuScenes-metadata 下载用于 mmdet 的元数据。

   否则

   请按照 MagicDrive-t 的方法将标注插值到 12Hz，并使用 tools/prepare_data/prepare_dataset.sh 中的命令自行生成元数据。

   如果你已经拥有来自 MagicDrive-t 的元数据文件，可以使用 tools/prepare_data/add_box_id.py 添加实例 ID 键。相关命令可在 tools/prepare_data/prepare_dataset.sh 中找到。

   你的数据文件夹应如下所示：

   ${CODE_ROOT}/data
   ├── nuscenes
   │   ├── ...
   │   └── interp_12Hz_trainval
   └── nuscenes_mmdet3d-12Hz
       ├── nuscenes_interp_12Hz_infos_train_with_bid.pkl
       └── nuscenes_interp_12Hz_infos_val_with_bid.pkl
   

3. （可选）为了加速数据加载，我们准备了 BEV 地图的 h5 格式缓存文件。

   说明

   这些缓存文件可以通过 tools/prepare_data/prepare_map_aux.py 使用 configs/cache_gen 中的不同配置生成。例如：

   python tools/prepare_data/prepare_map_aux.py +cache_gen=map_cache_gen_interp \
       +process=val +subfix=8x200x200_12Hz
   

   完整的命令可在 tools/prepare_data/prepare_dataset.sh 中找到。

   请确保将生成的缓存文件移动到正确的位置。我们的默认路径是：

   ${CODE_ROOT}/data/nuscenes_map_aux_12Hz
   ├── train_8x200x200_12Hz.h5 (25G)
   ├── train_8x400x400_12Hz.h5 (99G)
   ├── val_8x200x200_12Hz.h5 (5.3G)
   └── val_8x400x400_12Hz.h5 (22G)
   

尝试 MagicDrive-V2

大多数情况下，你可以在 GPU 服务器和 Ascend 服务器上使用相同的命令。

推理模型以进行生成

# ${GPUS} 可以是 1/2/4/8，用于序列并行。
# ${CFG} 可以是位于 `configs/magicdrive/inference/` 中的任何配置文件。
# ${PATH_TO_MODEL} 可以是 `ema.pt` 文件的路径，或检查点中 `model` 文件的路径。
# ${FRAME} 可以是 1/9/17/33/65/129/full...(8n+1)。1 表示单张图像；full 表示 nuScenes 的全长。
# 如果 GPU 显存充足，可以省略 `cpu_offload=true` 和 `scheduler.type=rflow-slice`。
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} num_frames=${FRAME} \
    cpu_offload=true scheduler.type=rflow-slice

有关 GPU 显存要求的更多信息，请参阅 FAQ。

例如，要以最高分辨率（848x1600）生成全长视频（20秒@12fps），使用 8 张 H20/A800 显卡：

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True torchrun --standalone --nproc_per_node 8 \
    scripts/inference_magicdrive.py \
    configs/magicdrive/inference/fullx848x1600_stdit3_CogVAE_boxTDS_wCT_xCE_wSST.py \
    --cfg-options model.from_pretrained=./ckpts/MagicDriveDiT-stage3-40k-ft/ema.pt \
    num_frames=full cpu_offload=true scheduler.type=rflow-slice

其他生成选项：

• force_daytime: (bool) 强制生成白天场景。
• force_rainy: (bool) 强制生成雨天场景。
• force_night: (bool) 强制生成夜晚场景。
• allow_class: (list) 限制生成的类别。
• del_box_ratio: (float) 随机删除部分框以进行生成。
• drop_nearest_car: (int) 在生成过程中删除距离最近的 N 辆车。

推理模型以进行测试

我们按照 W-CODA2024 Track2 的格式生成视频，并使用已建立的基准进行测试。在生成之前，请确保评估用的元数据已准备好，如下所示：

${CODE_ROOT}/data/nuscenes_mmdet3d-12Hz
├── nuscenes_interp_12Hz_infos_track2_eval.pkl # 可从 Track2 页面下载
└── nuscenes_interp_12Hz_infos_track2_eval_with_bid.pkl  # 可由本项目生成或下载。

使用 8 张 GPU/NPU 生成视频：

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True  # 对于 GPU
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True  # 对于 NPU
torchrun --standalone --nproc_per_node 8 scripts/test_magicdrive.py \
    configs/magicdrive/test/17-16x848x1600_stdit3_CogVAE_boxTDS_wCT_xCE_wSST_map0_fsp8_cfg2.0.py \
    --cfg-options model.from_pretrained=${PATH_TO_MODEL} tag=${TAG}

训练 MagicDrive-V2

使用 32 张 A800/H20 显卡启动训练：

# 请将“xx”替换为实际的 rank 和 ip
# ${config} 可以是 `configs/magicdrive/train` 中的任意配置文件。
# 例如：configs/magicdrive/train/stage3_higher-b-v3.1-12Hz_stdit3_CogVAE_boxTDS_wCT_xCE_wSST_bs4_lr1e-5_sp4simu8.py
torchrun --nproc-per-node=8 --nnode=4 --node_rank=xx --master_addr xx --master_port 18836 \
    scripts/train_magicdrive.py ${config} --cfg-options num_workers=2 prefetch_factor=2

我们还使用 64 张 Ascend 910b 来训练第 2 阶段，相关配置可在 configs/magicdrive/npu_64g 中找到。

此外，我们还提供了调试配置，用于测试你的环境和数据加载流程：

# 例如（使用 4 张 A800）
# ${config} 可以是 `configs/magicdrive/train` 中的任意配置文件。
# 例如：configs/magicdrive/train/stage3_higher-b-v3.1-12Hz_stdit3_CogVAE_boxTDS_wCT_xCE_wSST_bs4_lr1e-5_sp4simu8.py
bash scripts/launch_1node.sh 4 ${config} --cfg-options debug=true
	
# 通过设置 `vsdebug=true` 并使用 1 个进程，你可以使用 VSCode 的“附加模式”进行调试。

注意：sp=4（第 3 阶段）至少需要 4 张 GPU 才能运行。

引用我们

@inproceedings{gao2025magicdrive-v2,
  title={{MagicDrive-V2}: High-Resolution Long Video Generation for Autonomous Driving with Adaptive Control},
  author={Gao, Ruiyuan and Chen, Kai and Xiao, Bo and Hong, Lanqing and Li, Zhenguo and Xu, Qiang},
  booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
  year={2025}
}

致谢

我们采用了以下开源项目：

• BEVFusion：用于处理3D边界框和BEV地图的数据加载器
• Open-Sora：STDiT3模型及训练框架
• ColossalAI：用于并行计算和Zero2优化的框架
• CogVideoX：我们使用了他们的CogVAE

MagicDrive-V2 快速上手指南

MagicDrive-V2 是一个基于 DiT 架构的开源项目，旨在生成用于自动驾驶的高分辨率、长时序街景视频。它支持多样化的 3D 几何控制和多视角一致性，适用于自动驾驶仿真与数据增强。

1. 环境准备

系统要求

• 硬件: 代码已在 NVIDIA A800/H20 或 华为 Ascend 910b 服务器上测试通过。
• 操作系统: Linux (推荐 Ubuntu)
• Python: 建议 Python 3.8+

前置依赖

请确保已安装对应的 CUDA/CANN 驱动。本项目强烈建议使用 pip 直接配置环境，不推荐直接使用 conda 的 yaml 文件。

核心依赖版本参考

• NVIDIA 平台:
  • PyTorch == 2.4.0
  • Apex (需源码编译)
  • xformers >= 0.0.27
  • flash-attn >= 2.6.3
• Ascend (华为昇腾) 平台:
  • CANN 8.0RC2
  • PyTorch == 2.3.1 + torch-npu == 2.3.1
  • Apex (昇腾版)

2. 安装步骤

第一步：克隆代码库

git clone https://github.com/flymin/MagicDrive-V2.git
cd MagicDrive-V2

第二步：配置基础环境

根据您的硬件平台选择以下一种方式安装基础包。

方案 A：NVIDIA GPU (A800/H20 等)

1. 安装 PyTorch 及相关加速库（请根据实际 CUDA 版本调整）：

   pip install torch==2.4.0 torchvision==0.19.0 --index-url https://download.pytorch.org/whl/cu118
   # 编译安装 apex
   git clone https://github.com/NVIDIA/apex
   cd apex && pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
   cd ..
   # 安装其他加速库
   pip install xformers>=0.0.27 flash-attn>=2.6.3
   

   国内加速提示: 如下载缓慢，可添加 -i https://pypi.tuna.tsinghua.edu.cn/simple 参数。

2. 安装定制版 ColossalAI：

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

3. 安装其余依赖：

   pip install -r requirements/requirements.txt
   

方案 B：Ascend NPU (910b)

1. 参照华为官方文档配置 PyTorch-NPU 环境 (基于 CANN 8.0RC2)。
2. 安装定制版 ColossalAI (已移除 bitsandbytes 依赖)：

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

3. 安装其余依赖：

   pip install -r requirements/requirements.txt
   

第三步：下载预训练模型

请在项目根目录下创建文件夹并整理模型结构：

mkdir -p pretrained ckpts

1. VAE 模型: 下载 THUDM/CogVideoX-2b 中的 vae 子文件夹。
2. Text Encoder: 下载 google/t5-v1_1-xxl。
3. MagicDrive-V2 主模型: 下载 Stage-3 检查点 flymin/MagicDriveDiT-stage3-40k-ft。

最终目录结构应如下：

${CODE_ROOT}/
├── pretrained/
│   ├── CogVideoX-2b/
│   │   └── vae/
│   └── t5-v1_1-xxl/
└── ckpts/
    └── MagicDriveDiT-stage3-40k-ft/

提示: 国内用户可使用 Hugging Face 镜像站 (如 hf-mirror.com) 加速下载。

第四步：准备数据 (nuScenes)

1. 从 nuScenes 官网 下载数据集至 ./data/ 目录。
2. 从 HuggingFace 下载处理好的元数据 (mmdet 格式)。
3. 将元数据放入 data/nuscenes_mmdet3d-12Hz/ 目录。

(可选) 为加速数据加载，可运行 tools/prepare_data/prepare_map_aux.py 生成 BEV 地图缓存文件 (.h5)。

3. 基本使用

视频生成推理

以下命令展示了如何使用 8 张显卡生成高分辨率长视频（20 秒，12fps，分辨率 848x1600）。

注意: 如果显存充足，可省略 cpu_offload=true 和 scheduler.type=rflow-slice 参数。

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

torchrun --standalone --nproc_per_node 8 \
    scripts/inference_magicdrive.py \
    configs/magicdrive/inference/fullx848x1600_stdit3_CogVAE_boxTDS_wCT_xCE_wSST.py \
    --cfg-options model.from_pretrained=./ckpts/MagicDriveDiT-stage3-40k-ft/ema.pt \
    num_frames=full cpu_offload=true scheduler.type=rflow-slice

常用控制参数:

• num_frames: 设置帧数，支持 1 (单图), 9, 17, 33, 65, 129 或 full (完整序列)。
• force_daytime / force_night / force_rainy: 强制生成特定天气或时间场景。
• allow_class: 限制生成的物体类别列表。
• del_box_ratio: 随机丢弃部分检测框以增加多样性。

模型测试 (Benchmark)

若需按照 W-CODA2024 Track2 标准进行评估，请确保已准备好评估专用的元数据文件，然后运行：

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

torchrun --standalone --nproc_per_node 8 scripts/test_magicdrive.py \
    configs/magicdrive/test/17-16x848x1600_stdit3_CogVAE_boxTDS_wCT_xCE_wSST_map0_fsp8_cfg2.0.py \
    --cfg-options model.from_pretrained=${PATH_TO_MODEL}

(请将 ${PATH_TO_MODEL} 替换为实际的模型路径)