DreamWaltz

<div align="center"> <img src="https://yellow-cdn.veclightyear.com/ab5030c0/4913628d-f726-445a-9413-1a4bf5a93eb7.png" width="30%"> </div> <h1 align="center">💃DreamWaltz: 用复杂的3D可动画化头像创造场景</h1> <p align="center">

项目主页 | 论文 | arXiv | 海报

本仓库包含NeurIPS 2023论文的官方实现：

DreamWaltz: 用复杂的3D可动画化头像创造场景 <br>黄钰坤^1,2, 王佳楠¹, 曾爱玲¹, 曹赫¹, 戚贤彪¹, 石宇凯¹, 查正军², 张磊¹<br> ¹国际数字经济研究院   ²中国科学技术大学

新闻

• 2024/09/01: 感谢黄泽桓提供的DreamWaltz的threestudio实现！
• 2023/11/10: 训练和推理代码已发布。

简介

DreamWaltz是一个使用预训练的2D扩散模型ControlNet和人体参数模型SMPL进行文本驱动3D可动画化头像创建的学习框架。核心思想是从骨架条件的扩散监督中优化可变形的NeRF表示，确保3D一致性并推广到任意姿势。

<p align="middle"> <img src="https://yellow-cdn.veclightyear.com/ab5030c0/3e0da3c4-7bc7-40fd-941e-1ddb590d9317.gif" width="80%"> <br> <em>图1. DreamWaltz可以生成可动画化头像(a)并构建复杂场景(b)(c)(d)。</em> </p>

安装

本代码主要基于优秀的latent-nerf和stable-dreamfusion项目。请安装依赖：

pip install -r requirements.txt

instant-ngp的CUDA扩展在运行时构建，与stable-dreamfusion相同。

准备SMPL权重

我们使用smpl和vposer模型进行头像创建和动画学习，请按照smplx和human_body_prior中的说明下载模型权重，并构建具有以下结构的目录：

smpl_models
├── smpl
│   ├── SMPL_FEMALE.pkl
│   └── SMPL_MALE.pkl
│   └── SMPL_NEUTRAL.pkl
└── vposer
    └── v2.0
        ├── snapshots
        ├── V02_05.yaml
        └── V02_05.log

然后，在configs/paths.py中更新模型路径SMPL_ROOT和VPOSER_ROOT。

准备动作序列

您可能需要准备SMPL格式的人体动作序列来为生成的头像制作动画。我们的代码为AIST++提供了数据API，这是一个带有SMPL注释的高质量舞蹈视频数据库。请从此网站下载SMPL注释，构建具有以下结构的目录：

aist
├── gWA_sFM_cAll_d26_mWA5_ch13.pkl
├── gWA_sFM_cAll_d27_mWA0_ch15.pkl
├── gWA_sFM_cAll_d27_mWA2_ch17.pkl
└── ...

并在configs/paths.py中更新数据路径AIST_ROOT。

开始使用

DreamWaltz主要包括两个训练阶段：(I) 标准化虚拟形象创建和(II) 可动画虚拟形象学习。

以下命令也在run.sh中提供。

1. SMPL引导的NeRF初始化

使用从标准姿态SMPL网格渲染的蒙版图像预训练NeRF：

python train.py \
  --log.exp_name "pretrained" \
  --log.pretrain_only True \
  --prompt.scene canonical-A \
  --prompt.smpl_prompt depth \
  --optim.iters 10000

获得的预训练检查点可用于不同的文本提示。

2. 标准化虚拟形象创建

使用基于ControlNet的SDS学习基于NeRF的标准化虚拟形象表示：

text="一个木制机器人"
avatar_name="wooden_robot"
pretrained_ckpt="./outputs/pretrained/checkpoints/step_010000.pth"
# 预训练检查点可用于不同的文本提示

python train.py \
  --guide.text "${text}" \
  --log.exp_name "canonical/${avatar_name}" \
  --optim.ckpt "${pretrained_ckpt}" \
  --optim.iters 30000 \
  --prompt.scene canonical-A

3. 可动画虚拟形象学习

使用基于ControlNet的SDS学习基于NeRF的可动画虚拟形象表示：

text="一个木制机器人"
avatar_name="wooden_robot"
canonical_ckpt="./outputs/canonical/${avatar_name}/checkpoints/step_030000.pth"

python train.py \
  --animation True \
  --guide.text "${text}" \
  --log.exp_name "animatable/${avatar_name}" \
  --optim.ckpt "${canonical_ckpt}" \
  --optim.iters 50000 \
  --prompt.scene random \
  --render.cuda_ray False

4. 制作舞蹈视频

基于训练好的可动画虚拟形象表示和目标动作序列制作舞蹈视频：

scene="gWA_sFM_cAll_d27_mWA2_ch17,180-280"
# "gWA_sFM_cAll_d27_mWA2_ch17"是AIST++中动作序列的文件名
# "180-280"是视频帧索引范围：[180, 280]

avatar_name="wooden_robot"
animatable_ckpt="./outputs/animatable/${avatar_name}/checkpoints/step_050000.pth"

python train.py \
    --animation True \
    --log.eval_only True \
    --log.exp_name "videos/${avatar_name}" \
    --optim.ckpt "${animatable_ckpt}" \
    --prompt.scene "${scene}" \
    --render.cuda_ray False \
    --render.eval_fix_camera True

生成的视频可以在PROJECT_ROOT/outputs/videos/${avatar_name}/results/128x128/中找到。

结果

标准化虚拟形象

<p align="middle"> <image src="assets/canonical_half.gif" width="80%"> <br> <em>图2. DreamWaltz可以根据文本描述创建标准化虚拟形象。</em> </p>

可动画虚拟形象

<p align="middle"> <image src="assets/animation_sp.gif" width="80%"> <br> <em>图3. DreamWaltz可以根据给定的动作序列为标准化虚拟形象制作动画。</em> </p>

复杂场景

<p align="middle"> <image src="assets/animation_obj.gif" width="80%"> <br> <em>图4. DreamWaltz可以制作包含虚拟形象与物体交互的复杂3D场景。</em> </p> <p align="middle"> <image src="assets/animation_scene.gif" width="80%"> <br> <em>图5. DreamWaltz可以制作包含虚拟形象与场景交互的复杂3D场景。</em> </p> <p align="middle"> <image src="assets/animation_mp.gif" width="80%"> <br> <em>图6. DreamWaltz可以制作包含虚拟形象之间交互的复杂3D场景。</em> </p>

参考

如果您发现这个代码库对您的工作有用，请考虑按以下方式引用：

@inproceedings{huang2023dreamwaltz,
  title={{DreamWaltz: Make a Scene with Complex 3D Animatable Avatars}},
  author={Yukun Huang and Jianan Wang and Ailing Zeng and He Cao and Xianbiao Qi and Yukai Shi and Zheng-Jun Zha and Lei Zhang},
  booktitle={Advances in Neural Information Processing Systems},
  year={2023}
}

@inproceedings{huang2024dreamtime,
  title={{DreamTime: An Improved Optimization Strategy for Diffusion-Guided 3D Generation}},
  author={Yukun Huang and Jianan Wang and Yukai Shi and Boshi Tang and Xianbiao Qi and Lei Zhang},
  booktitle={International Conference on Learning Representations},
  year={2024}
}