OctFusion: Octree-based Diffusion Models for 3D Shape Generation

[arXiv] [BibTex]

Code release for the paper "OctFusion: Octree-based Diffusion Models for 3D Shape Generation".

1. Installation

1. Clone this repository

git clone https://github.com/octree-nn/octfusion.git
cd octfusion

2. Create a Conda environment.

conda create -n octfusion python=3.11 -y && conda activate octfusion

3. Install PyTorch with Conda

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

4. Install other requirements.

pip3 install -r requirements.txt 

2. Generation with pre-trained models

2.1 Download pre-trained models

We provide the pretrained models for the category-conditioned generation and sketch-conditioned generation. Please download the pretrained models from Google Drive or Baidu Netdisk and put them in saved_ckpt/diffusion-ckpt and saved_ckpt/vae-ckpt.

2.2 Generation

1. Unconditional generation in category airplane, car, chair, rifle, table.

sh scripts/run_snet_uncond.sh generate hr $category
# Example
sh scripts/run_snet_uncond.sh generate hr airplane

2. Category-conditioned generation

sh scripts/run_snet_cond.sh generate hr $category
# Example
sh scripts/run_snet_cond.sh generate hr chair

3. Train from scratch

3.1 Data Preparation

1. Download ShapeNetCore.v1.zip (31G) from ShapeNet and place it in data/ShapeNet/ShapeNetCore.v1.zip. Download filelist from Google Drive or Baidu Netdisk and place it in data/ShapeNet/filelist.

2. Convert the meshes in ShapeNetCore.v1 to signed distance fields (SDFs). We use the same data preparation as DualOctreeGNN. Note that this process is relatively slow, it may take several days to finish converting all the meshes from ShapeNet.

python tools/repair_mesh.py --run convert_mesh_to_sdf
python tools/repair_mesh.py --run generate_dataset

3.2 Train OctFusion

1. VAE Training. We provide pretrained weights in saved_ckpt/vae-ckpt/vae-shapenet-depth-8.pth.

sh scripts/run_snet_vae.sh train vae im_5

2. Train the first stage model. We provide pretrained weights in saved_ckpt/diffusion-ckpt/$category/df_steps-split.pth.

sh scripts/run_snet_uncond.sh train lr $category

3. Load the pretrained first stage model and train the second stage. We provide pretrained weights in saved_ckpt/diffusion-ckpt/$category/df_steps-union.pth.

sh scripts/run_snet_uncond.sh train hr $category

Citation

If you find this code helpful, please consider citing:

1. arxiv version

@article{xiong2024octfusion,
  author = {Xiong, Bojun and Wei, Si-Tong and Zheng, Xin-Yang and Cao, Yan-Pei and Lian, Zhouhui and Wang, Peng-Shuai},
  title = {{OctFusion}: Octree-based Diffusion Models for 3D Shape Generation},
  journal = {arXiv},
  year = {2024},
}

Issues and FAQ

Coming soon!

Acknowledgement

This code borrows heavely from SDFusion, LAS-Diffusion, DualOctreeGNN. We thank the authors for their great work. The followings packages are required to compute the SDF: mesh2sdf.