TreeMeshGPT

Follow these steps to set up the environment and install dependencies:

1. Clone and Environment Setup:

   git clone https://github.com/sail-sg/TreeMeshGPT.git
   cd TreeMeshGPT
   conda create -n tmgpt python=3.11
   conda activate tmgpt

2. Install PyTorch (Requires version $\ge$ 2.5.0 for FlexAttention support during training):

   pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu124

3. Install Dependencies:

   pip install -r requirements.txt

Download the pre-trained model checkpoints and place them in a checkpoints directory. There are two versions available:

• 7-bit model: Supports meshes with up to 5,500 faces.
• 9-bit model: Supports meshes with up to 11,000 faces.

Use the following commands to download and organize them:

mkdir checkpoints

# 7-bit model
gdown 1IERM76szBQq9oAMoFw1Sbgp3kx97
mv treemeshgpt_7bit.pt checkpoints/

# 9-bit model
gdown 19Auv48x7kgoODRS7dij8QijWzDqRjq97
mv treemeshgpt_9bit.pt checkpoints/

Training involves two main steps: preparing the data (tokenization) and running the training loop.

1. Prepare Data (Tokenization)

Run train_create_pkl.py to process meshes and create .pkl files for the dataloader. This script applies data augmentation and tokenization.

# Install wandb first
pip install wandb

# Run tokenization
python train_create_pkl.py

Note: For large datasets, consider running multiple instances of this script in parallel on different subsets of data.

2. Launch Training

Use accelerate to launch the training script. The current implementation is optimized for a batch size of 1.

accelerate launch --mixed_precision=fp16 train.py

Important Configuration: You must manually adjust MESH_IN and PKL_OUT in train_create_pkl.py, as well as TRAIN_PATH and VAL_PATH in train.py, to match your local data paths.

Generate an artistic mesh conditioned on a point cloud using the inference.py script. The output is saved to the generation folder.

Basic Usage

python inference.py

Common Scenarios

• Using the 9-bit model with a specific mesh:

  python inference.py --version 9bit --mesh_path demo/luma_bunny.glb

• Handling small triangles (Enable boundary deletion and lower target face count):

  python inference.py --version 9bit --mesh_path demo/luma_box.glb --decimation_boundary_deletion True --decimation_target_nfaces 2000

• Running without decimation (e.g., for Objaverse evaluation):

  python inference.py --decimation False --mesh_path demo/objaverse_pig.obj