DreamCraft3D

论文 | 项目页面 | Youtube视频 | 在线演示

DreamCraft3D：基于分层引导扩散模型的三维生成方法的官方实现

摘要：我们提出了DreamCraft3D，一种生成高保真和一致性三维对象的分层三维内容生成方法。我们通过利用二维参考图像来引导几何雕刻和纹理增强的阶段来解决这个问题。这项工作的核心重点是解决现有工作中遇到的一致性问题。为了雕刻具有连贯渲染的几何形状，我们通过视图依赖扩散模型进行分数蒸馏采样。这个三维先验结合了几种训练策略，优先考虑几何一致性，但牺牲了纹理保真度。我们进一步提出了引导的分数蒸馏来专门提升纹理。我们在场景的增强渲染上训练了一个个性化扩散模型Dreambooth，将其注入场景的三维知识。这个三维感知的扩散先验的分数蒸馏为场景提供了一致视图的引导。值得注意的是，通过交替优化扩散先验和三维场景表示，我们实现了相互强化的改进：优化后的三维场景有助于训练特定场景的扩散模型，后者提供越来越多的一致视图引导，为三维优化提供启发。这样，优化过程得以引导，从而显著提升纹理。在分层生成的整个过程中使用定制的三维先验，DreamCraft3D生成连贯且具有逼真渲染效果的三维对象，推动了三维内容生成的前沿发展。

方法概述

安装

安装threestudio

这一部分和原始的threestudio相同。如果您已经安装了环境，请跳过这一部分。

有关更多信息，包括通过Docker安装，请参阅installation.md。

您需要至少20GB VRAM的NVIDIA显卡并安装CUDA。
安装Python >= 3.8。
（可选，推荐）创建一个虚拟环境：

python3 -m virtualenv venv
. venv/bin/activate

# 新版本的pip，例如pip-23.x，比老版本（例如pip-20.x）要快得多。它会缓存git包的轮子以避免以后不必要的重建。
python3 -m pip install --upgrade pip

安装PyTorch >= 1.12。我们在torch1.12.1+cu113和torch2.0.0+cu118上进行过测试，但其他版本也应该可以正常工作。

# torch1.12.1+cu113
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
# 或者torch2.0.0+cu118
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

（可选，推荐）安装ninja以加快CUDA扩展的编译速度：

pip install ninja

安装依赖项：

pip install -r requirements.txt

下载预训练模型

Zero123。我们默认使用最新的stable-zero123.ckpt，可以在这里下载到load/zero123/目录中。在论文中我们使用的是zero123-xl.ckpt，可以通过以下方式下载：

cd load/zero123
bash download.sh

Omnidata。我们在preprocess_image.py中使用Omnidata来进行深度和法线预测（从stable-dreamfusion复制）。

cd load/omnidata
gdown '1Jrh-bRnJEjyMCS7f-WsaFlccfPjJPPHI&confirm=t' # omnidata_dpt_depth_v2.ckpt
gdown '1wNxVO4vVbDEMEpnAi_jwQObf2MFodcBR&confirm=t' # omnidata_dpt_normal_v2.ckpt

快速开始

预处理输入图像以移除背景并获得其深度和法线图像。

python preprocess_image.py /path/to/image.png --recenter

我们的模型分多个阶段训练。您可以通过以下指令运行它。

prompt="a brightly colored mushroom growing on a log"
image_path="load/images/mushroom_log_rgba.png"

# --------- 阶段 1 (NeRF & NeuS) --------- #
python launch.py --config configs/dreamcraft3d-coarse-nerf.yaml --train system.prompt_processor.prompt="$prompt" data.image_path="$image_path"

ckpt=outputs/dreamcraft3d-coarse-nerf/$prompt@LAST/ckpts/last.ckpt
python launch.py --config configs/dreamcraft3d-coarse-neus.yaml --train system.prompt_processor.prompt="$prompt" data.image_path="$image_path" system.weights="$ckpt"

# --------- 阶段 2 (几何细化) --------- #
ckpt=outputs/dreamcraft3d-coarse-neus/$prompt@LAST/ckpts/last.ckpt
python launch.py --config configs/dreamcraft3d-geometry.yaml --train system.prompt_processor.prompt="$prompt" data.image_path="$image_path" system.geometry_convert_from="$ckpt"


# --------- 阶段 3 (纹理细化) --------- #
ckpt=outputs/dreamcraft3d-geometry/$prompt@LAST/ckpts/last.ckpt
python launch.py --config configs/dreamcraft3d-texture.yaml --train system.prompt_processor.prompt="$prompt" data.image_path="$image_path" system.geometry_convert_from="$ckpt"

<details> <summary>[可选] 如果在第1阶段出现"Janus问题"，请考虑训练一个自定义的Text2Image模型。</summary>

首先，通过Zero123++从单个参考图像生成多视图图像。

python threestudio/scripts/img_to_mv.py --image_path 'load/mushroom.png' --save_path '.cache/temp' --prompt 'a photo of mushroom' --superres

通过DreamBooth Lora训练个性化的DeepFloyd模型。请检查上述生成的多视图图像是否合理。

export MODEL_NAME="DeepFloyd/IF-I-XL-v1.0"
export INSTANCE_DIR=".cache/temp"
export OUTPUT_DIR=".cache/if_dreambooth_mushroom"

accelerate launch threestudio/scripts/train_dreambooth_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a sks mushroom" \
  --resolution=64 \
  --train_batch_size=4 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --scale_lr \
  --max_train_steps=1200 \
  --checkpointing_steps=600 \
  --pre_compute_text_embeddings \
  --tokenizer_max_length=77 \
  --text_encoder_use_attention_mask

个性化的DeepFloyd模型Lora会保存在.cache/if_dreambooth_mushroom中。现在您可以通过以下指令替换训练脚本中的引导：

# --------- 阶段 1 (NeRF & NeuS) --------- #
python launch.py --config configs/dreamcraft3d-coarse-nerf.yaml --train system.prompt_processor.prompt="$prompt" data.image_path="$image_path" system.guidance.lora_weights_path=".cache/if_dreambooth_mushroom"

</details>

提示

内存使用。我们在40G A100 GPU上运行默认配置。为了减少内存使用，您可以通过data.height=128 data.width=128 data.random_camera.height=128 data.random_camera.width=128减小NeuS的渲染分辨率。您也可以通过相同方式减少其他阶段的分辨率。

导出网格

您可以使用以下脚本导出带纹理的网格obj文件：

# 这将使用默认的mesh-exporter配置导出obj+mtl
python launch.py --config path/to/trial/dir/configs/parsed.yaml --export --gpu 0 resume=path/to/trial/dir/ckpts/last.ckpt system.exporter_type=mesh-exporter

<details> <summary> 示例OBJ文件的可视化。 </summary> <img src="https://yellow-cdn.veclightyear.com/35dd4d3f/acbb5229-2786-4d49-915c-d0b4b02f8de7.jpg"> </details>

待办事项

发布重新整理的代码。
发布测试图像数据。
清理原始的dreambooth训练代码。
提供一些运行结果和检查点。

致谢

这个代码基于令人惊叹的开源项目threestudio-project和stable-dreamfusion构建。

BibTeX

@article{sun2023dreamcraft3d,
  title={Dreamcraft3d: Hierarchical 3d generation with bootstrapped diffusion prior},
  author={Sun, Jingxiang and Zhang, Bo and Shao, Ruizhi and Wang, Lizhen and Liu, Wen and Xie, Zhenda and Liu, Yebin},
  journal={arXiv preprint arXiv:2310.16818},
  year={2023}
}