EDGE — 官方 PyTorch 实现

EDGE：基于音乐的可编辑舞蹈生成（CVPR 2023）
Jonathan Tseng, Rodrigo Castellon, C. Karen Liu
https://arxiv.org/abs/2211.10658 摘要：舞蹈是人类重要的艺术形式，但创作新舞蹈往往困难且耗时。在本工作中，我们提出了可编辑舞蹈生成（EDGE），这是一种最先进的方法，能够在忠实于输入音乐的同时，生成逼真且物理上合理的舞蹈。EDGE 使用基于 Transformer 的扩散模型，并结合强大的音乐特征提取器 Jukebox，赋予了舞蹈生成强大的编辑能力，包括关节级别的条件控制和中间帧插值等。我们引入了一种新的物理合理性度量指标，并通过以下方式对所提出方法生成的舞蹈质量进行了全面评估：(1) 多项关于物理合理性、节拍对齐和多样性的定量指标；更重要的是，(2) 一项大规模用户研究，结果表明我们的方法显著优于先前的最先进方法。

需求

• 出于性能和兼容性考虑，我们建议使用 Linux 系统。Windows 系统可能也能运行，但未被官方支持。
• 64 位 Python 3.7+
• PyTorch 1.12.1
• 每张 GPU 至少需要 16 GB 内存
• 1–8 张高端 NVIDIA GPU，每张至少配备 16 GB 显存，安装 NVIDIA 驱动及 CUDA 11.6 工具包。

本仓库示例构建经过验证的环境：

• Debian 10
• 64 位 Python 3.7.12
• PyTorch 1.12.1
• 16 GB 内存
• 1 张 NVIDIA T4，CUDA 11.6 工具包

此外，本仓库还依赖以下库，这些库可能需要特殊的安装步骤：

• jukemirlib
• pytorch3d
• accelerate
  • 注意：安装后请务必运行 accelerate config。我们使用 fp16 精度。
• wine（可选，仅用于导入 Blender）

快速入门

快速开始

• 从 Google Drive 下载保存的模型检查点，或直接运行 bash download_model.sh。
• 运行 demo.ipynb，该笔记本演示了模型的基本接口。

加载自定义音乐

您可以通过将自定义音乐下载为 .wav 文件并放入指定目录（例如 custom_music/），然后运行以下命令来测试模型：

python test.py --music_dir custom_music/

此过程可能需要一些时间，因为脚本会将指定音乐的所有 Jukebox 特征提取并存储在内存中。这些特征也可以保存下来并在后续使用中重复利用，只需添加 --cache_features 和 --use_cached_features 参数即可。更多细节请参阅 args.py 文件。
注意：请确保文件名格式规范，例如 Britney Spears - Toxic (Official HD Video).wav 可能因空格和括号导致不可预测的行为，而 toxic.wav 则会按预期工作。演示笔记本中使用了 youtube-dl --output 标志来实现这一点。

（可选，仅限微调）数据集下载

使用以下命令下载并处理 AIST++ 数据集（仅包含音频和动作数据）：

cd data
bash download_dataset.sh
python create_dataset.py --extract-baseline --extract-jukebox

这将按照论文中的设置对数据集进行预处理。数据处理大约需要 24 小时，并占用约 50 GB 空间，以预先计算数据集中所有 Jukebox 特征。

训练您自己的模型

在 AIST++ 数据集下载并处理完成后，您可以运行训练脚本，例如：

accelerate launch train.py --batch_size 128  --epochs 2000 --feature_type jukebox --learning_rate 0.0002

以论文中设定的参数训练模型。训练过程会将进度记录到 wandb 并间歇性地生成样本输出以可视化学习效果。根据可用的 GPU 数量，整个训练过程可能需要 6 到 24 小时。

评估您的模型

使用论文中提出的物理足部接触（PFC）分数评估您的模型输出：

1. 生成约 1000 个样本，并使用 --save_motions 参数保存关节位置。
2. 运行评估脚本：

python test.py --music_dir custom_music/ --save_motions
python eval/eval_pfc.py

Blender 3D 渲染

为了以 3D 方式渲染生成的舞蹈，我们将它们转换为 FBX 文件，以便在 Blender 中使用。我们提供了一个示例骨骼绑定文件 SMPL-to-FBX/ybot.fbx。 在启用 --save-motions 标志生成舞蹈后，将相关的 .pkl 文件移动到一个文件夹中，例如 smpl_samples。 运行以下命令将动作转换为 FBX 文件，这些文件可以导入 Blender 并重新绑定到不同的骨骼上，例如 Mixamo 提供的骨骼。此外，还有多种重定向工具可供选择，例如 Rokoko 插件 for Blender。

开发说明

这是一个研究性质的实现，通常在发布后不会定期更新或维护。

引用

@article{tseng2022edge,
  title={EDGE: Editable Dance Generation From Music},
  author={Tseng, Jonathan and Castellon, Rodrigo and Liu, C Karen},
  journal={arXiv preprint arXiv:2211.10658},
  year={2022}
}

致谢

我们感谢 lucidrains 提供的 Adan 和 diffusion 代码库，感谢 softcat477 提供的 SMPL to FBX 库，以及 BobbyAnguelov 提供的 FBX 转换工具。

EDGE 快速上手指南

EDGE (Editable Dance Generation From Music) 是一个基于 PyTorch 的开源项目，能够根据输入音乐生成逼真且符合物理规律的舞蹈动作，并支持关节级条件控制和中间帧生成等编辑功能。

环境准备

系统要求

• 操作系统: 推荐 Linux (如 Debian 10)。Windows 可能可用但非官方支持。
• Python: 64-bit Python 3.7+ (推荐 3.7.12)。
• 内存: 每块 GPU 至少需要 16 GB RAM。
• GPU: 1–8 张高端 NVIDIA 显卡，显存至少 16 GB，需安装 NVIDIA 驱动和 CUDA 11.6 toolkit。
• 框架: PyTorch 1.12.1。

前置依赖

本项目依赖以下特殊库，安装前请确保已配置好基础环境：

• jukemirlib: 音乐特征提取器。
• pytorch3d: Facebook Research 开发的 3D 深度学习库。
• accelerate: Hugging Face 提供的加速库。
• wine (可选): 仅当需要将结果导入 Blender 时需要。

注意: 国内用户安装 pytorch3d 和 jukemirlib 时可能会遇到网络问题，建议检查是否有对应的国内镜像源或手动下载源码包进行本地安装。

安装步骤

1. 克隆仓库并安装基础依赖

   git clone <repository_url>
   cd EDGE
   pip install -r requirements.txt
   

   (注：若根目录无 requirements.txt，请根据项目实际依赖列表手动安装 torch==1.12.1 等核心库)

2. 安装特殊依赖库 按照各库官方说明安装 jukemirlib 和 pytorch3d。例如：

   pip install jukemirlib
   # pytorch3d 通常需要编译安装，请参考其 GitHub 页面针对 CUDA 11.6 的安装指令
   pip install accelerate
   

3. 配置 Accelerate 安装完成后，必须运行配置命令。根据项目要求，混合精度选择 fp16：

   accelerate config
   

4. 下载预训练模型 从 Google Drive 下载 checkpoint，或直接运行脚本：

   bash download_model.sh
   

   (国内用户若无法访问 Google Drive，需自行寻找模型文件 1BAR712cVEqB8GR37fcEihRV_xOC-fZrZ 并放置于对应目录)

基本使用

1. 体验演示 (Quickstart)

运行 Jupyter Notebook 查看模型的基本接口和示例：

jupyter notebook demo.ipynb

2. 使用自定义音乐生成舞蹈

将你的音乐文件转换为 .wav 格式放入目录（例如 custom_music/）。 重要提示: 文件名必须规范化，避免空格和括号。

• ❌ 错误示例: Britney Spears - Toxic (Official HD Video).wav
• ✅ 正确示例: toxic.wav

运行生成脚本：

python test.py --music_dir custom_music/

首次运行会提取 Jukebox 音乐特征，耗时较长。可通过 --cache_features 保存特征，后续使用 --use_cached_features 加速。

3. (可选) 导出为 3D 模型 (FBX)

若需在 Blender 中渲染，需先将生成的动作转换为 FBX 格式。

1. 生成动作时带上 --save_motions 参数，将生成的 .pkl 文件移至 SMPL-to-FBX/smpl_samples/ 目录。
2. 执行转换脚本：

   python SMPL-to-FBX/Convert.py --input_dir SMPL-to-FBX/smpl_samples/ --output_dir SMPL-to-FBX/fbx_out
   

3. 将输出的 FBX 文件导入 Blender，可使用 Mixamo 或 Rokoko 插件进行骨骼重定向。