DisCo: 解构控制的现实人类舞蹈生成

Tan Wang*, Linjie Li*, Kevin Lin*, Yuanhao Zhai, Chung-Ching Lin, Zhengyuan Yang, Hanwang Zhang, Zicheng Liu, Lijuan Wang

南洋理工大学 | 微软 Azure AI | 布法罗大学

:fire: 新闻

[2024.07.15] 我们推出了 IDOL，这是 DisCo 的增强版，可同时生成视频和深度，实现现实的 2.5D 视频合成。
[2024.04.08] DisCo 已被 CVPR24 接收。请查看最新版本的 ArXiv 论文。
[2023.12.30] 更新的幻灯片介绍 DisCo 并总结最近的工作。
[2023.11.30] 更新了带时间模块的 DisCo。
[2023.10.12] 更新了新的 ArXiv 版本的 DisCo（增加了时间模块；与 MCVD 同步 FVD 计算；更多基线和可视化等）。
[2023.07.21] 更新了 TSV 文件的构建指南。
[2023.07.08] 更新了 Colab 演示（确保我们的代码/演示可以在任何机器上运行）！
[2023.07.03] 提供了本地演示部署示例代码。现在您可以在自己的开发机器上尝试我们的演示！
[2023.07.03] 我们更新了预训练 tsv 数据。
[2023.06.28] 我们发布了 DisCo 人类属性预训练代码。
[2023.06.21] 发布了 DisCo 人类图像编辑演示！试试看！
[2023.06.21] 我们发布了人类特定微调代码供参考。来构建你自己的特定舞蹈模型！
[2023.06.21] 发布了解构控制的一般微调代码。
[2023.06.21] 发布了人类属性预训练的检查点和微调检查点。

其他您可能感兴趣的项目：

来自阿里巴巴的Animate Anyone
来自抖音的MagicAnimate
来自抖音的MagicDance

🎨 Gradio 演示

本地启动演示（视频舞蹈生成演示即将推出）

下载微调检查点模型（我们的演示使用这个检查点 checkpoint，您也可以使用自己的模型）；通过 git clone https://huggingface.co/lambdalabs/sd-image-variations-diffusers 下载 sd-image-variation。
运行 jupyter notebook 文件。所有必需的代码/命令已设置好。记得修改预训练模型路径 --pretrained_model 和 --pretrained_model_path (sd-va) 在 manual_args = [xxx]。
运行后，这个 jupyter 将自动使用您的本地开发 GPU 启动演示。您可以通过 notebook 末尾提供的 web 链接访问演示。
或者您可以参考我们在 Colab 上的部署。所有代码都是从头部署的！

[在线 Gradio 演示]（时间）

📝 介绍

在本项目中，我们介绍了 DisCo，它是一种广义的参考人类舞蹈生成工具包，支持 人类图像和视频生成 及 多种使用场景（预训练、微调和特定人类的微调），尤其适用于现实场景。

✨与现有工作相比，DisCo 实现了：

通用性：无需进行以人类为特定对象的微调即可用于大规模现实世界中的人类（我们也支持人类特定微调）。之前的方法仅支持生成特定领域的人类，例如，DreamPose仅能生成拥有简单双管步姿态的时尚模特。
当前针对人类舞蹈生成的SOTA结果。
广泛的使用案例和应用（更多详情请见项目页面）。
简单易懂的框架，支持高效训练（x-formers，FP16训练，deepspeed，wandb）以及多种可能的研究方向（预训练 -> 微调 -> 人类特定微调）。

🌟有了这个项目，你可以得到：

[用户]：只需尝试我们的在线演示，或在本地部署模型推理。
[研究人员]：一个易于使用的代码库用于重新实现和开发。
[研究人员]：大量的研究方向以便进一步改进。

🚀 快速入门

安装

## 在创建py3.8环境之后
pip install --user torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
pip install --user progressbar psutil pymongo simplejson yacs boto3 pyyaml ete3 easydict deprecated future django orderedset python-magic datasets h5py omegaconf einops ipdb
pip install --user --exists-action w -r requirements.txt
pip install git+https://github.com/microsoft/azfuse.git


## 加速
pip install --user deepspeed==0.6.3
pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers

## 你可能需要将protobuf降级至3.20.x版本
pip install protobuf==3.20.0

数据准备

1. 人类属性预训练

我们创建了一个人类图像子集（70万张图像）用于人类属性预训练，从现有的图像语料库中过滤得来：

数据集	COCO（单人）	TikTok风格	DeepFashion2	SHHQ-1.0	LAION-Human
大小	2万	12.4万	27.6万	4万	24万

经过预处理的有效TSV数据格式的预训练数据可以下载这里 (Google Drive)。

数据根目录
└── composite/
    ├── train_xxx.yaml  # 这个路径需要在训练参数中指定
    └── val_xxx.yaml
    ...
└── TikTokDance/
    ├── xxx_images.tsv
    └── xxx_poses.tsv
    ...
└── coco/  
    ├── xxx_images.tsv
    └── xxx_poses.tsv

2. 使用解耦控制进行微调

我们使用TikTok数据集进行微调。

我们已经预处理了TikTok数据并使用了高效的TSV格式，可下载这里 (Google Drive)。（请注意，我们只使用每个TikTok视频的第一帧作为参考图像。）

数据文件夹的结构应该如下：

数据根目录
└── composite_offset/
    ├── train_xxx.yaml  # 这个路径需要在训练参数中指定
    └── val_xxx.yaml
    ...
└── TikTokDance/
    ├── xxx_images.tsv
    └── xxx_poses.tsv
    ...

*PS: 如果你希望使用你自己的数据资源但采用我们的TSV数据结构，请参考PREPRO.MD。

人类属性预训练

训练：

AZFUSE_USE_FUSE=0 QD_USE_LINEIDX_8B=0 NCCL_ASYNC_ERROR_HANDLING=0 python finetune_sdm_yaml.py --cf config/ref_attn_clip_combine_controlnet_attr_pretraining/coco_S256_xformers_tsv_strongrand.py --do_train --root_dir /home1/wangtan/code/ms_internship2/github_repo/run_test \
--local_train_batch_size 64 --local_eval_batch_size 64 --log_dir exp/tiktok_pretrain \
--epochs 40 --deepspeed --eval_step 2000 --save_step 2000 --gradient_accumulate_steps 1 \
--learning_rate 1e-3 --fix_dist_seed --loss_target "noise" \
--train_yaml ./blob_dir/debug_output/video_sythesis/dataset/composite/train_TiktokDance-coco-single_person-Lindsey_0411_youtube-SHHQ-1.0-deepfashion2-laion_human-masks-single_cap.yaml --val_yaml ./blob_dir/debug_output/video_sythesis/dataset/composite/val_TiktokDance-coco-single_person-SHHQ-1.0-masks-single_cap.yaml \
--unet_unfreeze_type "transblocks" --refer_sdvae --ref_null_caption False --combine_clip_local --combine_use_mask \
--conds "masks" --max_eval_samples 2000 --strong_aug_stage1 --node_split_sampler 0

预训练模型检查点：OneDrive

使用解耦控制进行微调

1. 修改配置文件

从官方diffusers repo下载sd-image-variations-diffusers并根据配置文件pretrained_model_path进行放置。或者你也可以选择修改pretrained_model_path。

2. 无分类器自由指南（CFG）

训练：

[*要启用WANDB，请在 utils/lib.py中设置wandb键]

[*要使用多GPU运行，试着在python前添加mpirun -np {GPU NUM}。]

WANDB_ENABLE=0 AZFUSE_USE_FUSE=0 NCCL_ASYNC_ERROR_HANDLING=0 mpirun -np 8 python finetune_sdm_yaml.py --cf config/ref_attn_clip_combine_controlnet/tiktok_S256L16_xformers_tsv.py \
--do_train --root_dir /home1/wangtan/code/ms_internship2/github_repo/run_test \
--local_train_batch_size 32 \
--local_eval_batch_size 32 \
--log_dir exp/tiktok_ft \
--epochs 20 --deepspeed \
--eval_step 500 --save_step 500 \
--gradient_accumulate_steps 1 \
--learning_rate 2e-4 --fix_dist_seed --loss_target "noise" \
--train_yaml /home/wangtan/data/disco/yaml_file/train_TiktokDance-poses-masks.yaml \
--val_yaml /home/wangtan/data/disco/yaml_file/new10val_TiktokDance-poses-masks.yaml \
--unet_unfreeze_type "all" \
--guidance_scale 3 \
--refer_sdvae \
--ref_null_caption False \
--combine_clip_local --combine_use_mask \
--conds "poses" "masks" \
--stage1_pretrain_path /path/to/pretrained_model_checkpoint/mp_rank_00_model_states.pt

可视化：

要运行可视化，只需将--do_train改为--eval_visu。你还可以使用'--eval_save_filename' xxx来指定可视化文件夹名称。

评估：

您首先需要运行评估以获取结果。然后我们使用gen_eval.sh一站式获取{exp_dir_path}/{prediction_folder_name}的评估指标

bash gen_eval.sh {exp_dir_path} {exp_dir_path}/{prediction_folder_name}

例如，

bash gen_eval.sh /home/kevintw/code/disco/github2/DisCo/save_results/TikTok_cfg_check /home/kevintw/code/disco/github2/DisCo/save_results/TikTok_cfg_check/pred_gs1.5_scale-cond1.0-ref1.0/

您可能需要下载预训练的视觉模型并在gen_eval.sh中修改路径以获得fvd指标。

模型检查点（OneDrive）：TikTok 训练数据 | 更多 TikTok 风格训练数据

3. 使用无分类器引导（CFG） [CFG 可以带来稍微更好的结果]

训练（在没有 CFG 的训练脚本中添加以下参数）：

--drop_ref 0.05 # 训练期间丢弃参考图像的概率
--guidance_scale 1.5 # CFG 的比例

可视化：

要运行可视化，只需将 --do_train 更改为 --eval_visu。你也可以通过 '--eval_save_filename' xxx 指定可视化文件夹名称。（记得还要指定 --guidance_scale）

你也可以查看我们的命令 bash 文件 config/command_bash/tiktok_cfg.sh 以供参考。

评估：

同上

模型检查点（OneDrive）：TikTok 训练数据 | 更多 TikTok 风格训练数据

4. 时间模块微调

训练： 在训练完成图像 DisCo 模型后，我们进一步结合了时间卷积层和时间注意层来提高时间上的平滑性。请注意，参数 --pretrained_model 的内容应该是图像 DisCo 模型检查点，而不是阶段 1 的预训练检查点。

AZFUSE_USE_FUSE=0 NCCL_ASYNC_ERROR_HANDLING=0 python finetune_sdm_yaml.py --cf config/disco_w_tm/yz_tiktok_S256L16_xformers_tsv_temdisco_temp_attn.py \
--do_train --root_dir /home1/wangtan/code/ms_internship2/github_repo/run_test \
--local_train_batch_size 2 \
--local_eval_batch_size 2 \
--log_dir exp/tiktok_ft \
--epochs 20 --deepspeed \
--eval_step 500 --save_step 500 \
--gradient_accumulate_steps 1 \
--learning_rate 1e-4 --fix_dist_seed --loss_target "noise" \
--train_yaml /home/wangtan/data/disco/yaml_file/train_TiktokDance-poses-masks.yaml \
--val_yaml /home/wangtan/data/disco/yaml_file/new10val_TiktokDance-poses-masks.yaml \
--unet_unfreeze_type "all" \
--refer_sdvae \
--ref_null_caption False \
--combine_clip_local --combine_use_mask \
--train_sample_interval 4 \
--nframe 16 \
--frame_interval 1 \
--conds "poses" "masks" \
--pretrained_model /path/to/pretrained_model/mp_rank_00_model_states.pt

评估： 只需用 gen_eval_tm.sh 脚本替换先前的 gen_eval.sh 脚本，如下所示。GT 文件夹路径将自动填充。

bash gen_eval_tm.sh {exp_dir_path} {exp_dir_path}/{prediction_folder_name}

模型检查点：TikTok 训练数据（有 CFG） | 更多 TikTok 风格训练数据（有 CFG）

5. FVD 评价指标再现的潜在问题

请先查看 GitHub 问题及回复这里。我们已经在 A100 GPU 上验证了检查点结果。如果你仍然无法再现结果，请提出问题或发送电子邮件给我。

特定人类细化训练

1. 准备你想用于训练的数据集

准备一个特定人类的视频或一组人类图像
使用 Grounded-SAM 和 OpenPose 获取每个训练图像的人类掩码和骨架（参见 PREPRO.MD 获取更多详情）
对于特定人类的细化训练，我们建议直接使用原始图像/掩码/姿势进行训练，而不是构建 TSV 文件。如果你仍然想使用 TSV 文件结构准备您的数据，请参见 PREPRO.MD 参考。

2. 运行以下脚本进行特定人类的细化训练：

对于参数调优，建议首先调整 learning-rate 和 unet_unfreeze_type。

AZFUSE_USE_FUSE=0 NCCL_ASYNC_ERROR_HANDLING=0 python finetune_sdm_yaml.py \
--cf config/ref_attn_clip_combine_controlnet_imgspecific_ft/webtan_S256L16_xformers_upsquare.py --do_train --root_dir /path/of/saving/root \
--local_train_batch_size 32 --local_eval_batch_size 32 --log_dir exp/human_specific_ft/ \
--epochs 20 --deepspeed --eval_step 500 --save_step 500 --gradient_accumulate_steps 1 \
--learning_rate 1e-3  --fix_dist_seed  --loss_target "noise" \
--unet_unfreeze_type "crossattn" \
--refer_sdvae --ref_null_caption False --combine_clip_local --combine_use_mask --conds "poses" "masks" \
--freeze_pose True --freeze_background False \
--pretrained_model /path/to/the/ft_model_checkpoint \
--ft_iters 500 --ft_one_ref_image False --ft_idx dataset/folder/name --strong_aug_stage1 True --strong_rand_stage2 True

备注

1. FVD 评价指标再现的潜在问题

请先查看 GitHub 问题及回复这里。我们已经在 A100 GPU 上验证了检查点结果。如果你仍然无法再现结果，请提出问题或发送电子邮件给我。

2. PSNR 评价指标

我们感谢 @Delicious-Bitter-Melon 突出了 PSNR 评价指标实现中的潜在数值溢出问题。这个问题已经在代码库中解决了，更新的评分反映在我们的最新版本中。需要注意的是，这个修正不会改变趋势或影响总体结论。

3. DisCo 不仅限于上半身人类。

请检查我们最新的论文了解其普适性。它可以通过在训练阶段结合广泛的图像缩放增强来轻松实现。

发布计划

“解耦控制细化训练”代码
“特定人类细化训练”代码
预训练和细化训练的模型检查点
HuggingFace 演示
“人类属性预训练”代码

引用

如果你在研究中使用了我们的工作，请引用：

@article{wang2023disco,
  title={Disco: Disentangled control for realistic human dance generation},
  author={Wang, Tan and Li, Linjie and Lin, Kevin and Zhai, Yuanhao and Lin, Chung-Ching and Yang, Zhengyuan and Zhang, Hanwang and Liu, Zicheng and Wang, Lijuan},
  journal={arXiv preprint arXiv:2307.00040},
  year={2023}
}