normal-depth-diffusion

<p align="center"> <br> <img src="https://yellow-cdn.veclightyear.com/835a84d5/df2448f5-3089-4846-a04f-5313df42b040.gif" width="400"/> <br> <h1>法线深度扩散模型</h1> <p>

法线深度扩散模型：一个通用的法线深度扩散模型。

如果您熟悉中文，可以阅读中文版本的README。

文本到法线深度

teaser-nd

文本到多视角法线深度

项目页面 | 论文 | YouTube

新闻

2023-12-25：我们通过阿里云OSS服务发布了训练数据集mvs_objaverse。我们还提供了一个方便的多线程脚本用于快速下载。
2023-12-11：推理代码和预训练模型已发布。我们正在努力改进ND-Diffusion Model，敬请期待！

3D生成

本仓库仅包含RichDreamer论文中的扩散模型和2D图像生成代码。
对于3D生成，请查看RichDreamer。

推理准备

使用以下脚本安装依赖项。

conda create -n nd
conda activate nd 
pip install -r requirements.txt
pip install git+https://github.com/openai/CLIP.git
pip install git+https://github.com/CompVis/taming-transformers.git
pip install webdataset
pip install img2dataset

我们还提供了一个dockerfile来构建docker镜像。

sudo docker build -t mv3dengine_22.04:cu118 -f docker/Dockerfile .

下载预训练权重。

ND：在Laion-2B上训练的法线深度扩散模型
ND-MV：多视角法线深度扩散模型
Alebdo-MV：多视角深度条件反照率扩散模型

我们还提供了一个下载脚本。

python tools/download_models/download_nd_models.py

推理（采样）

我们提供了一个用于采样的脚本

sh demo_inference.sh

或使用以下详细说明：

文本到法线深度采样

# dmp求解器
python ./scripts/t2i.py --ckpt $ckpt_path --prompt $prompt --dpm_solver --n_samples 2 --save_dir $save_dir
# plms求解器
python ./scripts/t2i.py --ckpt $ckpt_path --prompt $prompt --plms --n_samples 2  --save_dir $save_dir
# ddim求解器
python ./scripts/t2i.py --ckpt $ckpt_path --prompt $prompt --n_samples 2  --save_dir $save_dir

文本到多视角法线深度采样

# nd-mv
python ./scripts/t2i_mv.py --ckpt_path $ckpt_path --prompt $prompt  --num_frames 4  --model_name nd-mv --save_dir $save_dir

# nd-mv带VAE（即将推出）
python ./scripts/t2i_mv.py --ckpt_path $ckpt_path --prompt $prompt  --num_frames 4  --model_name nd-mv-vae --save_dir $save_dir

文本到多视角反照率采样

python ./scripts/td2i_mv.py --ckpt_path $ckpt_path --prompt $prompt --depth_file $ depth_file --num_frames 4  --model_name albedo-mv --save_dir $save_dir

训练准备

下载Laion-2B-en-5-AES（训练ND模型所需）

从parquet下载laion-2b数据集然后，将parquet文件放入./laion2b-dataset-5-aes

cd ./tools/download_dataset
bash ./download_2b-5_aes.sh
cd -

下载单目先验模型权重（训练ND模型所需）

NormalBae scannet.pt
Midas3.1 dpt_beit_large512.pt

# 将scannet.pt移动到normalbae先验模型
mv scannet.pt ./libs/ControlNet-v1-1-nightly/annotator/normalbae/scannet.pt
# 将dpt_beit_large512.pt移动到./libs/omnidata_torch/pretrained_models/dpt_beit_large_512.pt
mv dpt_beit_large512.pt ./libs/omnidata_torch/pretrained_models/dpt_beit_large_512.pt

下载Objaverse数据集的渲染多视角图像（训练ND-MV和Albedo-MV模型所需）

使用准备好的脚本下载我们的渲染数据集

wget https://virutalbuy-public.oss-cn-hangzhou.aliyuncs.com/share/aigc3d/valid_paths_v4_cap_filter_thres_28.json
# 示例：python ./scripts/data/download_objaverse.py ./mvs_objaverse ./valid_paths_v4_cap_filter_thres_28.json 50
python ./scripts/data/download_objaverse.py /保存数据的路径 /valid_paths_v4_cap_filter_thres_28.json的路径 线程数(例如 10)
# 如果将数据保存在其他位置，请设置一个链接
ln -s /保存数据的路径 mvs_objaverse
# 说明文件
wget https://virutalbuy-public.oss-cn-hangzhou.aliyuncs.com/share/aigc3d/text_captions_cap3d.json

训练

训练法线深度VAE模型

下载在ImageNet上预训练的VAE权重。
修改configs/autoencoder_normal_depth/autoencoder_normal_depth.yaml中的配置文件，设置model.ckpt_path=/预训练VAE权重的路径

# 训练VAE数据集
bash ./scripts/train_vae/train_nd_vae/train_rgbd_vae_webdatasets.sh \ model.params.ckpt_path=${预训练VAE权重} \
data.params.train.params.curls='path_laion/{00000..${:5 end_id}}.tar' \
--gpus 0,1,2,3,4,5,6,7

训练法线深度扩散模型

在训练并获得法线深度VAE模型后，或者您可以从ND-VAE下载

# 第一步
export SD-MODEL-PATH=/sd-1.5的路径
bash scripts/train_normald_sd/txt_cond/web_datasets/train_normald_webdatasets.sh --gpus 0,1,2,3,4,5,6,7 \
    model.params.first_stage_ckpts=${法线深度VAE} model.params.ckpt_path=${SD-MODEL-PATH} \
    data.params.train.params.curls='path_laion/{00000..${:5 end_id}}.tar'

# 第二步 在./configs/stable-diffusion/normald/sd_1_5/txt_cond/web_datasets/laion_2b_step2.yaml中修改您的step_weights路径
bash scripts/train_normald_sd/txt_cond/web_datasets/train_normald_webdatasets_step2.sh --gpus 0,1,2,3,4,5,6,7 \
    model.params.first_stage_ckpts=${法线深度VAE} \
    model.params.ckpt_path=${预训练步骤权重} \
    data.params.train.params.curls='path_laion/{00000..${:5 end_id}}.tar'

训练多视角法线深度扩散模型

在训练并获得法线深度扩散模型后，或者您可以从ND下载，

我们提供两个版本的多视角法线深度扩散模型

a. 不带VAE去噪 b. 带VAE去噪

在当前版本中，我们提供不带VAE去噪的版本

# a. 训练不带VAE版本
bash ./scripts/train_normald_sd/txt_cond/objaverse/objaverse_finetune_wovae_mvsd-4.sh --gpus 0,1,2,3,4,5,6,7,  \
    model.params.ckpt_path=${法线深度扩散}
# b. 训练带VAE版本
bash ./scripts/train_normald_sd/txt_cond/objaverse/objaverse_finetune_mvsd-4.sh --gpus 0,1,2,3,4,5,6,7, \
    model.params.ckpt_path=${法线深度扩散}

训练多视角深度条件反照率扩散模型

在训练并获得法线深度扩散模型后，或者您可以从ND下载，

bash scripts/train_abledo/objaverse/objaverse_finetune_mvsd-4.sh --gpus 0,1,2,3,4,5,6,7, model.params.ckpt_path=${法线深度扩散}

致谢

我们大量借鉴了以下代码库。非常感谢这些作者分享他们的代码。

引用

@article{qiu2023richdreamer,
    title={RichDreamer: A Generalizable Normal-Depth Diffusion Model for Detail Richness in Text-to-3D}, 
    author={Lingteng Qiu and Guanying Chen and Xiaodong Gu and Qi zuo and Mutian Xu and Yushuang Wu and Weihao Yuan and Zilong Dong and Liefeng Bo and Xiaoguang Han},
    year={2023},
    journal = {arXiv preprint arXiv:2311.16918}
}