A unified framework that couples 3D assembly with 3D generation: assembly provides part-level structural priors for generation, while generation injects holistic shape context that resolves ambiguities in assembly — enabling pose prediction and missing-geometry synthesis in a single model.
Zeyu Jiang · Sihang Li · Siqi Tan · Chenyang Xu · Juexiao Zhang · Julia Galway-Witham · Xue Wang · Scott A. Williams · Radu Iovita · Chen Feng✉ · Jing Zhang✉
✉ Corresponding author
2026/05/25: We release the CRAG codebase. Pretrained checkpoints and data will be released shortly — stay tuned!2026/05/01: CRAG has been accepted to ICML 2026 🎉
Todos
- Release pretrained checkpoints (Stage 1 & Stage 2, Breaking Bad, PartNeXt and Omni variants)
We recommend using uv to manage the dependencies. Follow the instructions here to install uv. Then simply run
uv sync
source .venv/bin/activateThe project is pinned to Python 3.12.3 and PyTorch 2.8.0 + CUDA 12.8.
If you run into issues, try removing the .venv directory and running uv clean && uv self-update before re-installing.
CRAG trains and evaluates across multiple assembly datasets. All datasets are stored in HDF5 format aligned to GARF with additional schema fields for renderings. We provide our processed datasets below:
| Dataset | Source | Download |
|---|---|---|
| Breaking Bad | Project Page | OneDrive |
| PartNeXt | Project Page | OneDrive |
We provide an evaluation script in scripts/eval.sh, modify the experiment and ckpt_path arguments to point to your desired checkpoint and config.
Results (per-sample transformations, generated results and aggregate metrics) are written under logs/<EXPERIMENT_NAME>/. For the most reliable per-sample dumps we recommend single-GPU evaluation.
Training proceeds in two stages: an assembly pre-training stage, followed by a generation-coupled fine-tuning stage. We provide one launcher per dataset family.
Breaking Bad + MorphoSource + Skull:
bash scripts/train_bb.sh stg_1_bbPartNeXt:
bash scripts/train_partnext.sh stg_1_partnextBoth scripts auto-resume from output/<EXPERIMENT_NAME>/last.ckpt if it exists, copy the HDF5 files to /dev/shm for faster I/O, and launch via torchrun using the PET_* environment variables provided by your distributed launcher.
Stage 2 jointly fine-tunes the assembly backbone together with the generative branch (a TripoSG-based shape head):
bash scripts/train_bb.sh stg_2_bb_finetune
# or
bash scripts/train_partnext.sh stg_2_partnext_finetuneYou will typically want to point ckpt_path at the Stage 1 checkpoint the first time you launch Stage 2. The launcher's auto-resume logic will then take over for subsequent runs.
For frequently asked questions, please refer to our GitHub Issues page. Search existing threads, or open a new issue if your question hasn't been answered yet.
If you find this project useful, please consider citing our paper:
@inproceedings{jiang2026crag,
title = {CRAG: Can 3D Generative Models Help 3D Assembly?},
author = {Jiang, Zeyu and Li, Sihang and Tan, Siqi and Xu, Chenyang and Zhang, Juexiao and Galway-Witham, Julia and Wang, Xue and Williams, Scott A. and Iovita, Radu and Feng, Chen and Zhang, Jing},
booktitle = {International Conference on Machine Learning (ICML)},
year = {2026}
}Our codebase builds on the excellent work of GARF and TripoSG.
This project is licensed under the GPL License. See the LICENSE file for details.
We thank the authors of the Breaking Bad and PartNeXt datasets for releasing the data that underpins our pre-training and evaluation, and the GARF and TripoSG projects whose codebases CRAG is built on top of.