TimSchneider42/cod-vae-16x16-tiny

COD-VAE 16 x 16 (tiny)

A tiny, decode-first COD-VAE that compresses a 3D shape into 16 latent vectors of 16 dimensions = 256 numbers and decodes them back into an occupancy field. Same latent shape as cod-vae-16x16 and cod-vae-16x16-small, but built for pipelines whose wall clock is the decode forward+backward through a frozen decoder (e.g. reconstruction-reward RL): ~6.6M parameters, roughly 4x faster than -small and 33x faster than the full-size model.

Note: the latent shape matches its siblings, but every model defines its own latent space — latents from one cannot be decoded with another.

Trained with cod-vae, a PyTorch/JAX reimplementation of COD-VAE (Cho et al., ICCV 2025). The weights are a self-contained npz and load with either backend. The JAX decode numbers below include the channel-last plane layout (cod-vae >= 56b2c82).

Architecture vs the -small recipe

	cod-vae-16x16-small	this model
embed dim / heads	256 / 4	128 / 4
encoder	3 blocks x 3 layers, 512 patches, mlp 4	2 blocks x 2 layers, 256 patches, mlp 2
refinement decoder	6 layers, 16-px patches	4 layers, 32-px patches
query planes (query_dim)	16 channels at 128²	8 channels at 96²
latent decoder layers	12	6
total parameters	~35M	~6.6M

The shipped config pins attention_implementation="default" (the XLA path), which is measurably faster than cuDNN's fused kernel on these short sequences.

Decode speed (H100, JAX float16, batch 1024 x 2048 queries, fwd+bwd through the full latent)

num_latents and latent_dim barely move the decode cost, so these numbers (measured on the 16x8 variant) hold for the whole -tiny family.

model	step	throughput
cod-vae-16x8 (full)	~350 ms	2.9k shapes/s
cod-vae-16x8-small	43.5 ms	23.6k shapes/s
cod-vae-16x8-tiny	8.0 ms	127k shapes/s

Held-out reconstruction quality

source	held-out shapes	volume IoU	near-surface accuracy
ABC (CAD parts)	128	0.7997	0.7738

For reference, cod-vae-16x16-small reaches 0.872 / 0.830 on the same protocol — the extra ~4x decode speedup costs additional quality. The 16x8 configuration was qualified against a hard floor of 0.75 ABC volume IoU before the grid was trained.

Usage

import trimesh
from cod_vae import CODVAE

vae = CODVAE.from_pretrained("TimSchneider42/cod-vae-16x16-tiny")

mesh = trimesh.load("bunny.obj", force="mesh")
latent, transform = vae.encode_mesh(mesh, return_transform=True)   # (16, 16)
reconstruction = vae.decode_mesh(latent, transform=transform)      # trimesh.Trimesh

Latents can also be computed from raw surface point clouds and decoded at arbitrary query points:

latents = vae.encode(points)                          # (N, 3) in [-1, 1]^3
logits = vae.decode(latents, queries)                 # occupancy logits, positive inside
volume = vae.decode_volume(latents, resolution=128)   # dense logit grid

Install with pip install cod-vae[torch,hub] (or cod-vae[jax,hub]).

Training data and recipe

Same 110,077-shape merged dataset and two-stage recipe as the -small grid: a 200-epoch stage-1 trunk per num_latents (shared by its row) and a fresh 100-epoch stage 2 per cell with 6 latent-decoder layers; see the training guide for the exact commands.

Citation

@inproceedings{cho2025cod,
  author={Cho, In and Yoo, Youngbeom and Jeon, Subin and Kim, Seon Joo},
  title={Representing 3D Shapes with 64 Latent Vectors for 3D Diffusion Models},
  booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
  year={2025}
}