TimSchneider42 commited on
Commit
ecbea24
·
verified ·
1 Parent(s): 12a3974

tiny res96 recipe, 200-epoch trunk, stage-2 e99

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