DiffusionRig

DiffusionRig：学习个性化先验用于面部外观编辑<br> Zheng Ding、Xuaner Zhang、 Zhihao Xia、Lars Jebe、 Zhuowen Tu、Xiuming Zhang <br>CVPR 2023<br> arXiv / 项目主页 / 视频 / BibTex

设置与准备

环境设置

conda create -n diffusionrig python=3.8
conda activate diffusionrig
conda install pytorch=1.11 cudatoolkit=11.3 torchvision -c pytorch
conda install mpi4py dlib scikit-learn scikit-image tqdm -c conda-forge
pip install lmdb opencv-python kornia yacs blobfile chumpy face_alignment

您还需要安装pytorch3d来渲染物理缓冲区：

conda install -c fvcore -c iopath -c conda-forge fvcore iopath
pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py38_cu113_pyt1110/download.html

DECA设置

在准备训练数据之前，请先下载DECA的源文件和检查点来设置它（您需要创建一个账户来下载FLAME资源）：

1. deca_model.tar：访问此页面下载预训练的DECA模型。
2. generic_model.pkl：访问此页面下载FLAME 2020并提取generic_model.pkl。
3. FLAME_texture.npz：访问同一页面下载FLAME texture space并提取FLAME_texture.npz。
4. 从DECA的数据页面下载以下列出的其他文件，并将它们也放在data/文件夹中：

data/
  deca_model.tar
  generic_model.pkl
  FLAME_texture.npz
  fixed_displacement_256.npy
  head_template.obj
  landmark_embedding.npy
  mean_texture.jpg
  texture_data_256.npy
  uv_face_eye_mask.png
  uv_face_mask.png

数据准备

我们使用FFHQ来训练第一阶段，使用个人相册来训练第二阶段。在训练之前，您需要使用DECA提取这些图像的物理缓冲区。

对于FFHQ，您需要先对图像进行对齐：

python scripts/create_data.py --data_dir PATH_TO_FFHQ_ALIGNED_IMAGES --output_dir ffhq256_deca.lmdb --image_size 256 --use_meanshape False

对于个人相册（在我们的实验中每个身份使用大约20张照片），将所有图像放入一个文件夹，然后运行以下命令进行对齐：

python scripts/align.py -i PATH_TO_PERSONAL_PHOTO_ALBUM -o personal_images_aligned -s 256

然后，通过运行以下命令创建数据集：

python scripts/create_data.py --data_dir personal_images_aligned --output_dir personal_deca.lmdb --image_size 256 --use_meanshape True

训练

第一阶段：学习通用人脸先验

我们的256x256模型在第一阶段训练中使用八个GPU，每个GPU的批量大小为32：

mpiexec -n 8 python scripts/train.py --latent_dim 64 --encoder_type resnet18 \
    --log_dir log/stage1 --data_dir ffhq256_deca.lmdb --lr 1e-4 \
    --p2_weight True --image_size 256 --batch_size 32 --max_steps 50000 \
    --num_workers 8 --save_interval 5000 --stage 1

要使模型无限期地继续训练，请设置--max_steps 0。如果您想恢复训练过程，只需添加--resume_checkpoint PATH_TO_THE_MODEL。

:white_check_mark: 我们还提供了我们训练的第一阶段模型在这里，以便您可以快速进入训练个性化模型的阶段。

第二阶段：学习个性化先验

在您的小型个人相册上微调模型：

mpiexec -n 1 python scripts/train.py --latent_dim 64 --encoder_type resnet18 \
    --log_dir log/stage2 --resume_checkpoint log/stage1/[MODEL_NAME].pt \
    --data_dir peronsal_deca.lmdb --lr 1e-5 \
    --p2_weight True --image_size 256 --batch_size 4 --max_steps 5000 \
    --num_workers 8 --save_interval 5000 --stage 2

在单个Nvidia V100 GPU上大约需要30分钟。

(2023年10月13日更新)：我们发布了一些我们在论文中使用的第二阶段训练数据集，仅供研究使用。请通过Google Drive链接下载它们。

推理

我们提供了一个脚本来通过修改物理缓冲区来编辑人脸外观。运行：

python scripts/inference.py --source SOURCE_IMAGE_FILE --target TARGET_IMAGE_FILE --output_dir OUTPUT_DIR --modes light --model_path PATH_TO_MODEL --meanshape PATH_TO_MEANSHAPE --timestep_respacing ddim20 

使用目标图像的物理参数（例如，光照、表情或头部姿势）来编辑源图像。

问题或疑问？

如果问题与代码相关，请在此处提出issue。

对于问题，也请考虑提出issue，因为这可能会对未来的读者有帮助。否则，请发送电子邮件至 zhding@ucsd.edu联系Zheng Ding。

致谢

这个代码库基于Guided-Diffusion、DECA和Diff-AE，并从中获得了灵感。我们感谢这些作者公开了这些代码库。