README.md

---
license: mit
pipeline_tag: text-to-3d
language:
- en
tags:
- 3d
- text-to-3d
- voxel
- diffusion
- rectified-flow
- custom-code
- tiny
---

<table align="center">
<tr>
<td align="center">

<a href="https://huggingface.co/spaces/bench-labs/blog?post=Vm1.html">📝 Blog Post</a>

</td>
</tr>
</table>

# VoxelModel v1 🧊

<img src="examples.png" alt="Twelve 32x32x32 grids generated from text prompts" style="width: 640px; max-width: 100%;">

A tiny text-to-3D diffusion model. You give it a prompt, it gives you a
32x32x32 occupancy grid. It is 40,050,892 trained parameters and it trained
in under ten hours on a single A100.

This is the PixelModel line stepping out of 2D. v1 through v4 were tiny latent
diffusion for images; VoxelModel keeps the same argument (low resolution as a
deliberate aesthetic, not a compromise) and applies it to shape. `model.png` is
still here and still holds every weight, because that gimmick ports unchanged to
3D: it only ever cared about floats.

## The one honest result

It works for bulky objects and it fails for thin ones.

Chairs, benches, tables, mushrooms, trees, cars and boats come back legible. A
sword comes back as 0.2% occupancy, which is a handful of stray voxels, and a key
comes back as a featureless bar.

That is not a training failure. A sword does not survive voxelization at 32^3 in
the **training data** either: the real Objaverse sword we voxelized for reference
sits at 2.9% occupancy and reads as a cross at best. The model reproduced what it
was shown. If you want swords, the answer is resolution, not more steps.

So: 32^3 is enough, for a specific and checkable definition of enough. That was
the question v1 existed to answer.

## No VAE

VoxelModel operates directly on voxels, so there is no autoencoder in the stack.
Total footprint is 40.05M trained plus 37.8M frozen CLIP ViT-B/32, which is
77.9M. PixelModel v4 was 161M for comparable trained size, because it carries an
83.7M frozen `sd-vae-ft-mse`. This one is tiny without an asterisk.

## Architecture

`VoxelDiT` is a DiT with 3D patchify. A 32^3 grid at patch 4 gives 512 tokens,
dim 384, depth 12, heads 6. Conditioning is cross-attention on frozen CLIP
ViT-B/32 text features (40 tokens), plus adaLN-zero from the pooled embedding and
the timestep. Positional encoding is a 3D sincos over the 8x8x8 patch grid.

Objective is rectified flow with `target = x1 - x0` and logit-normal timestep
sampling, `t = sigmoid(randn)`.

## Training

| | |
|---|---|
| meshes | 28,415 train, 1,024 held out |
| steps | 90,000 at batch 256 |
| hardware | one A100 80GB, 2.57 steps/s, 41.4 GiB |
| wall clock | 9.7 hours |
| optimizer | AdamW 2e-4, betas (0.9, 0.99), grad clip 1.0 |
| schedule | cosine to 1e-6, 500 step warmup |
| EMA | 0.9999 |
| CFG dropout | 10% |

Held-out loss bottomed at **0.018737 at step 64,000** and drifted up to 0.019314
by step 90,000. That is 3.08% of overfitting, and it means the run was about 2.8
hours longer than it needed to be. 90,000 steps over 28,415 meshes is 811 epochs
with no augmentation, which is the obvious thing to fix next. Yaw rotations plus
a mirror would give 8x for free on a cubic grid.

The shipped weights are the step 90,000 EMA, not the step 64,000 best, because
checkpoints were written to a fixed path and overwrote each other. Noted, and
worth 3.08% to you if you retrain.

## Data

Meshes are `dylanebert/objaverse-lowpoly-obj`, captions are `tiange/Cap3D`,
joined on Objaverse UID for 58,374 captioned meshes. Voxelization centers on the
vertex mean, scales by max absolute coordinate, and rasterizes surfaces at pitch
2/31 with no fill, since filling solidifies closed volumes and pushes 4% of the
set past 90% occupancy.

`caption_filter.py` then throws away half of it. Multi-object meshes are the main
defect in this corpus: normalization scales the whole scene bounding box into the
grid, so a mesh holding several objects shrinks each one into a cluster of blobs.
The caption gives this away before you load anything, so the filter rejects scene
and variant-set wording, a top level `and` joining two noun phrases, and captions
over 12 words. Colour pairs survive on purpose, so "a purple and green mushroom"
is kept and "a white cube and a white robot" is not.

Measured over all 58,374 captions, 29,439 survive, which is 50.4%.

Worth knowing if you build on this: **occupancy does not detect the defect.**
Survivors and rejects have almost the same occupancy distribution, medians 4.3%
against 5.0% over a 900 mesh sample. No threshold would have caught it.

## Usage

```bash
python sample.py "a wooden chair" "a purple mushroom" --cfg 5.0 --steps 50
```

Writes a JSON of base64 `np.packbits` grids. To load the weights from the PNG
instead of the safetensors:

```python
from png_codec import load_model_png
model = load_model_png("model.png", "model_png.json", device="cuda")
```

Round-trip error against the safetensors is 9.7e-04 max absolute, which is fp16
quantization and nothing else.

## What v1 does not do

**Colour.** These meshes carry no vertex colours and no `.mtl` files, so the
model generates grey geometry only. Colour needs a GLB-with-materials source and
its own voxelization path.

**A real eval metric.** FID does not transfer to 3D. Held-out rectified flow loss
is what the table above reports, and it is a training diagnostic, not a quality
measure. Rendered-FID with CLIP score on 2D renders, or IoU against a held-out
split, are the candidates. This is unresolved and we would rather say so than
publish a number we do not believe.

## Files

| file | what |
|---|---|
| `model.safetensors` | EMA weights, fp32 |
| `model.png` | the same weights as a 6329x6329 PNG |
| `model_png.json` | parameter manifest for the PNG |
| `voxel_dit.py` | the model |
| `train.py` | training loop |
| `sample.py` | inference |
| `caption_filter.py` | the single-object filter |
| `build_data.py` | mesh download and voxelization |
| `png_codec.py` | PNG encode, decode and verify |