tiny res96 recipe, 200-epoch trunk, stage-2 e99
Browse files
README.md
ADDED
|
@@ -0,0 +1,108 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: mit
|
| 3 |
+
library_name: cod-vae
|
| 4 |
+
pipeline_tag: feature-extraction
|
| 5 |
+
tags:
|
| 6 |
+
- 3d
|
| 7 |
+
- shape-reconstruction
|
| 8 |
+
- autoencoder
|
| 9 |
+
- vae
|
| 10 |
+
- occupancy
|
| 11 |
+
---
|
| 12 |
+
|
| 13 |
+
# COD-VAE 16 x 8 (tiny)
|
| 14 |
+
|
| 15 |
+
A tiny, decode-first [COD-VAE](https://arxiv.org/abs/2503.08737) that compresses a
|
| 16 |
+
3D shape into **16 latent vectors of 8 dimensions = 128 numbers** and
|
| 17 |
+
decodes them back into an occupancy field. Same latent shape as
|
| 18 |
+
[cod-vae-16x8](https://huggingface.co/TimSchneider42/cod-vae-16x8) and
|
| 19 |
+
[cod-vae-16x8-small](https://huggingface.co/TimSchneider42/cod-vae-16x8-small),
|
| 20 |
+
but built for pipelines whose wall clock is the decode forward+backward through a
|
| 21 |
+
frozen decoder (e.g. reconstruction-reward RL): ~6.6M parameters, roughly
|
| 22 |
+
**4x faster than `-small` and 33x faster than the full-size model**.
|
| 23 |
+
|
| 24 |
+
> **Note:** the latent shape matches its siblings, but every model defines its own
|
| 25 |
+
> latent space — latents from one cannot be decoded with another.
|
| 26 |
+
|
| 27 |
+
Trained with [`cod-vae`](https://github.com/TimSchneider42/cod-vae), a PyTorch/JAX
|
| 28 |
+
reimplementation of COD-VAE (Cho et al., ICCV 2025). The weights are a
|
| 29 |
+
self-contained npz and load with either backend. The JAX decode numbers below
|
| 30 |
+
include the channel-last plane layout (cod-vae >= 56b2c82).
|
| 31 |
+
|
| 32 |
+
## Architecture vs the -small recipe
|
| 33 |
+
|
| 34 |
+
| | cod-vae-16x8-small | this model |
|
| 35 |
+
|---|---|---|
|
| 36 |
+
| embed dim / heads | 256 / 4 | 128 / 4 |
|
| 37 |
+
| encoder | 3 blocks x 3 layers, 512 patches, mlp 4 | 2 blocks x 2 layers, 256 patches, mlp 2 |
|
| 38 |
+
| refinement decoder | 6 layers, 16-px patches | 4 layers, 32-px patches |
|
| 39 |
+
| query planes (`query_dim`) | 16 channels at 128² | 8 channels at 96² |
|
| 40 |
+
| latent decoder layers | 12 | 6 |
|
| 41 |
+
| total parameters | ~35M | ~6.6M |
|
| 42 |
+
|
| 43 |
+
The shipped config pins `attention_implementation="default"` (the XLA path), which
|
| 44 |
+
is measurably faster than cuDNN's fused kernel on these short sequences.
|
| 45 |
+
|
| 46 |
+
## Decode speed (H100, JAX float16, batch 1024 x 2048 queries, fwd+bwd through the full latent)
|
| 47 |
+
|
| 48 |
+
`num_latents` and `latent_dim` barely move the decode cost, so these numbers
|
| 49 |
+
(measured on the 16x8 variant) hold for the whole `-tiny` family.
|
| 50 |
+
|
| 51 |
+
| model | step | throughput |
|
| 52 |
+
|---|---|---|
|
| 53 |
+
| cod-vae-16x8 (full) | ~350 ms | 2.9k shapes/s |
|
| 54 |
+
| cod-vae-16x8-small | 43.5 ms | 23.6k shapes/s |
|
| 55 |
+
| **cod-vae-16x8-tiny** | **8.0 ms** | **127k shapes/s** |
|
| 56 |
+
|
| 57 |
+
## Held-out reconstruction quality
|
| 58 |
+
|
| 59 |
+
| source | held-out shapes | volume IoU | near-surface accuracy |
|
| 60 |
+
|---|---|---|---|
|
| 61 |
+
| ABC (CAD parts) | 128 | 0.7665 | 0.7497 |
|
| 62 |
+
|
| 63 |
+
For reference, cod-vae-16x8-small reaches 0.842 / 0.804 on the same protocol — the extra
|
| 64 |
+
~4x decode speedup costs additional quality. The 16x8 configuration was qualified
|
| 65 |
+
against a hard floor of 0.75 ABC volume IoU before the grid was trained.
|
| 66 |
+
|
| 67 |
+
## Usage
|
| 68 |
+
|
| 69 |
+
```python
|
| 70 |
+
import trimesh
|
| 71 |
+
from cod_vae import CODVAE
|
| 72 |
+
|
| 73 |
+
vae = CODVAE.from_pretrained("TimSchneider42/cod-vae-16x8-tiny")
|
| 74 |
+
|
| 75 |
+
mesh = trimesh.load("bunny.obj", force="mesh")
|
| 76 |
+
latent, transform = vae.encode_mesh(mesh, return_transform=True) # (16, 8)
|
| 77 |
+
reconstruction = vae.decode_mesh(latent, transform=transform) # trimesh.Trimesh
|
| 78 |
+
```
|
| 79 |
+
|
| 80 |
+
Latents can also be computed from raw surface point clouds and decoded at arbitrary
|
| 81 |
+
query points:
|
| 82 |
+
|
| 83 |
+
```python
|
| 84 |
+
latents = vae.encode(points) # (N, 3) in [-1, 1]^3
|
| 85 |
+
logits = vae.decode(latents, queries) # occupancy logits, positive inside
|
| 86 |
+
volume = vae.decode_volume(latents, resolution=128) # dense logit grid
|
| 87 |
+
```
|
| 88 |
+
|
| 89 |
+
Install with `pip install cod-vae[torch,hub]` (or `cod-vae[jax,hub]`).
|
| 90 |
+
|
| 91 |
+
## Training data and recipe
|
| 92 |
+
|
| 93 |
+
Same 110,077-shape merged dataset and two-stage recipe as the `-small` grid: a
|
| 94 |
+
200-epoch stage-1 trunk per `num_latents` (shared by its row) and a fresh 100-epoch
|
| 95 |
+
stage 2 per cell with 6 latent-decoder layers; see the
|
| 96 |
+
[training guide](https://github.com/TimSchneider42/cod-vae/blob/main/TRAINING.md)
|
| 97 |
+
for the exact commands.
|
| 98 |
+
|
| 99 |
+
## Citation
|
| 100 |
+
|
| 101 |
+
```bibtex
|
| 102 |
+
@inproceedings{cho2025cod,
|
| 103 |
+
author={Cho, In and Yoo, Youngbeom and Jeon, Subin and Kim, Seon Joo},
|
| 104 |
+
title={Representing 3D Shapes with 64 Latent Vectors for 3D Diffusion Models},
|
| 105 |
+
booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
|
| 106 |
+
year={2025}
|
| 107 |
+
}
|
| 108 |
+
```
|