VoxelModel v1: tiny text-to-3D voxel diffusion

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,149 @@
+---
+license: mit
+pipeline_tag: text-to-3d
+language:
+- en
+tags:
+- 3d
+- text-to-3d
+- voxel
+- diffusion
+- rectified-flow
+- custom-code
+- tiny
+---
+
+# 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 cost
+$3.01 to train.
+
+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, $3.01 |
+| 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 |

build_data.py

@@ -0,0 +1,61 @@
+import json, os, sys, time, numpy as np, trimesh
+from multiprocessing import Pool
+sys.path.insert(0, '/root')
+from caption_filter import is_single_object
+
+root = '/root/lowpoly'
+OUT = '/root/data'
+os.makedirs(OUT, exist_ok=True)
+
+paths = {}
+for dirpath, _, names in os.walk(root):
+    for n in names:
+        if n.endswith('.obj'):
+            paths[n[:-4]] = os.path.join(dirpath, n)
+print('obj files on disk:', len(paths), flush=True)
+
+caps = json.load(open('/root/captions.json'))
+uids = sorted(u for u in caps if u in paths)
+print('captioned meshes:', len(uids), flush=True)
+
+def vox(u):
+    try:
+        m = trimesh.load(paths[u], force='mesh')
+        v = m.vertices - m.vertices.mean(0)
+        s = np.abs(v).max()
+        m.vertices = v / (s if s > 0 else 1)
+        g = m.voxelized(pitch=2.0/31).matrix.astype(np.uint8)
+        pad = np.zeros((32,32,32), np.uint8)
+        sh = [min(32,d) for d in g.shape]
+        pad[:sh[0],:sh[1],:sh[2]] = g[:sh[0],:sh[1],:sh[2]]
+        if pad.sum() == 0:
+            return u, None
+        return u, np.packbits(pad.reshape(-1))
+    except Exception:
+        return u, None
+
+t0 = time.time()
+ok_uids, rows, failed = [], [], 0
+with Pool(32) as p:
+    for i, (u, packed) in enumerate(p.imap(vox, uids, chunksize=64)):
+        if packed is None:
+            failed += 1
+        else:
+            ok_uids.append(u); rows.append(packed)
+        if (i+1) % 5000 == 0:
+            print(f'  voxelized {i+1}/{len(uids)}  failed {failed}  {time.time()-t0:.0f}s', flush=True)
+
+vox_arr = np.stack(rows)
+np.save(f'{OUT}/voxels_packed.npy', vox_arr)
+captions = [caps[u] for u in ok_uids]
+keep = [bool(is_single_object(c)) for c in captions]
+json.dump({'uids': ok_uids, 'captions': captions, 'keep': keep},
+          open(f'{OUT}/meta.json', 'w'))
+
+occ = np.unpackbits(vox_arr, axis=1).mean(1) * 100
+print(f'\ndone in {time.time()-t0:.0f}s')
+print(f'meshes voxelized {len(ok_uids)}, failed/empty {failed}')
+print(f'array {vox_arr.shape} {vox_arr.nbytes/2**20:.0f} MiB packed')
+print(f'filter survivors {sum(keep)} = {sum(keep)/len(keep)*100:.1f}%')
+print(f'surface occupancy mean {occ.mean():.2f}%  median {np.median(occ):.2f}%')
+print('BUILDDONE', flush=True)

caption_filter.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import re
+
+MAX_WORDS = 12
+
+SCENE_PATTERNS = [
+    ("variation", r"\bvari(?:ation|ations|ant|ants|ous)\b"),
+    ("set of", r"\bsets?\s+of\b"),
+    ("collection", r"\bcollections?\b"),
+    ("assorted", r"\bassorted\b"),
+    ("multiple", r"\bmultiple\b"),
+    ("pair of", r"\bpairs?\s+of\b"),
+    ("group of", r"\bgroups?\s+of\b"),
+    ("scene", r"\bscenes?\b"),
+    ("elements", r"\belements?\b"),
+    ("featuring", r"\bfeaturing\b"),
+    ("different", r"\bdifferent\b"),
+]
+
+MODIFIERS = {
+    "black", "white", "red", "blue", "green", "yellow", "orange", "purple",
+    "pink", "brown", "gray", "grey", "gold", "golden", "silver", "bronze",
+    "copper", "beige", "tan", "teal", "cyan", "magenta", "maroon", "navy",
+    "ivory", "cream", "turquoise", "lavender", "olive", "crimson", "scarlet",
+    "dark", "light", "pale", "bright", "deep", "pastel", "neon", "striped",
+    "wooden", "wood", "metal", "metallic", "plastic", "stone", "glass",
+    "steel", "iron", "brick", "concrete", "leather", "fabric", "marble",
+}
+
+_WORD = re.compile(r"[a-z][a-z'-]*")
+_AND = re.compile(r"\band\b")
+
+def _joins_noun_phrases(text):
+    parts = _AND.split(text)
+    if len(parts) < 2:
+        return False
+    for left, right in zip(parts[:-1], parts[1:]):
+        if "," in left:
+            return True
+        lw = _WORD.findall(left)
+        rw = _WORD.findall(right)
+        if not lw or not rw:
+            return True
+        if lw[-1] in MODIFIERS and rw[0] in MODIFIERS:
+            continue
+        return True
+    return False
+
+def reject_reason(caption):
+    if not caption or not caption.strip():
+        return "empty"
+    text = caption.lower()
+    if len(caption.split()) > MAX_WORDS:
+        return "too long"
+    for name, pattern in SCENE_PATTERNS:
+        if re.search(pattern, text):
+            return name
+    if _joins_noun_phrases(text):
+        return "and"
+    return None
+
+def is_single_object(caption):
+    return reject_reason(caption) is None

config.json

@@ -0,0 +1,14 @@
+{
+  "dit": {
+    "vox_size": 32,
+    "patch": 4,
+    "dim": 384,
+    "depth": 12,
+    "heads": 6,
+    "text_dim": 512
+  },
+  "steps": 90000,
+  "batch_size": 256,
+  "meshes": 28415,
+  "filtered": true
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8cce2211e7e7500dd661df15e64e185770046879073394e5a4a31a0185331f0a
+size 161012360

model_png.json

@@ -0,0 +1 @@
+{"cfg": {"vox_size": 32, "patch": 4, "dim": 384, "depth": 12, "heads": 6, "text_dim": 512}, "params": [{"name": "x_embed.weight", "shape": [384, 64], "numel": 24576}, {"name": "x_embed.bias", "shape": [384], "numel": 384}, {"name": "t_mlp.0.weight", "shape": [384, 384], "numel": 147456}, {"name": "t_mlp.0.bias", "shape": [384], "numel": 384}, {"name": "t_mlp.2.weight", "shape": [384, 384], "numel": 147456}, {"name": "t_mlp.2.bias", "shape": [384], "numel": 384}, {"name": "text_proj.weight", "shape": [384, 512], "numel": 196608}, {"name": "text_proj.bias", "shape": [384], "numel": 384}, {"name": "text_pool.weight", "shape": [384, 512], "numel": 196608}, {"name": "text_pool.bias", "shape": [384], "numel": 384}, {"name": "blocks.0.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.0.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.0.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.0.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.0.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.0.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.0.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.0.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.0.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.0.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.0.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.0.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.0.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.0.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.0.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.0.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.0.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.0.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.0.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.1.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.1.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.1.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.1.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.1.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.1.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.1.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.1.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.1.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.1.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.1.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.1.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.1.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.1.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.1.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.1.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.1.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.1.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.1.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.2.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.2.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.2.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.2.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.2.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.2.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.2.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.2.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.2.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.2.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.2.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.2.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.2.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.2.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.2.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.2.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.2.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.2.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.2.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.3.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.3.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.3.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.3.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.3.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.3.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.3.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.3.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.3.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.3.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.3.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.3.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.3.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.3.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.3.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.3.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.3.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.3.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.3.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.4.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.4.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.4.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.4.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.4.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.4.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.4.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.4.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.4.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.4.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.4.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.4.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.4.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.4.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.4.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.4.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.4.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.4.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.4.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.5.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.5.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.5.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.5.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.5.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.5.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.5.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.5.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.5.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.5.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.5.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.5.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.5.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.5.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.5.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.5.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.5.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.5.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.5.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.6.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.6.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.6.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.6.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.6.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.6.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.6.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.6.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.6.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.6.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.6.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.6.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.6.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.6.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.6.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.6.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.6.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.6.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.6.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.7.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.7.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.7.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.7.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.7.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.7.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.7.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.7.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.7.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.7.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.7.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.7.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.7.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.7.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.7.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.7.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.7.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.7.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.7.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.8.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.8.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.8.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.8.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.8.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.8.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.8.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.8.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.8.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.8.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.8.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.8.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.8.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.8.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.8.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.8.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.8.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.8.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.8.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.9.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.9.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.9.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.9.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.9.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.9.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.9.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.9.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.9.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.9.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.9.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.9.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.9.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.9.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.9.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.9.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.9.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.9.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.9.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.10.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.10.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.10.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.10.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.10.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.10.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.10.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.10.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.10.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.10.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.10.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.10.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.10.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.10.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.10.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.10.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.10.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.10.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.10.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.11.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.11.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.11.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.11.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.11.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.11.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.11.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.11.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.11.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.11.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.11.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.11.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.11.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.11.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.11.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.11.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.11.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.11.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.11.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "ada_out.1.weight", "shape": [768, 384], "numel": 294912}, {"name": "ada_out.1.bias", "shape": [768], "numel": 768}, {"name": "head.weight", "shape": [64, 384], "numel": 24576}, {"name": "head.bias", "shape": [64], "numel": 64}], "total_parameters": 40050892, "side": 6329, "dtype": "float16", "channels": "R=hi,G=lo,B=unused"}

png_codec.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+
+import numpy as np
+import torch
+from PIL import Image
+from safetensors.torch import load_file
+
+from voxel_dit import VoxelDiT
+
+def encode(weights, png_path, cfg_path, config="config.json"):
+    c = json.load(open(config))["dit"]
+    model = VoxelDiT(vox_size=c["vox_size"], patch=c["patch"], dim=c["dim"],
+                     depth=c["depth"], heads=c["heads"], text_dim=c["text_dim"])
+    model.load_state_dict(load_file(weights))
+    parts, manifest = [], []
+    for name, p in model.named_parameters():
+        a = p.detach().to(torch.float16).contiguous().view(-1).numpy()
+        parts.append(a)
+        manifest.append({"name": name, "shape": list(p.shape), "numel": int(a.size)})
+    flat = np.concatenate(parts)
+    N = flat.size
+    side = math.ceil(math.sqrt(N))
+    u16 = flat.view(np.uint16)
+    img = np.zeros((side * side, 3), dtype=np.uint8)
+    img[:N, 0] = (u16 >> 8).astype(np.uint8)
+    img[:N, 1] = (u16 & 0xFF).astype(np.uint8)
+    Image.fromarray(img.reshape(side, side, 3), "RGB").save(png_path)
+    total = sum(m["numel"] for m in manifest)
+    json.dump({"cfg": c, "params": manifest, "total_parameters": total,
+               "side": side, "dtype": "float16", "channels": "R=hi,G=lo,B=unused"},
+              open(cfg_path, "w"))
+    mb = os.path.getsize(png_path) / 1e6
+    print(f"[png] encoded {total:,} params -> {side}x{side} PNG ({mb:.1f} MB)")
+    return total, side
+
+def load_model_png(png_path, cfg_path, device="cpu"):
+    meta = json.load(open(cfg_path))
+    c = meta["cfg"]
+    model = VoxelDiT(vox_size=c["vox_size"], patch=c["patch"], dim=c["dim"],
+                     depth=c["depth"], heads=c["heads"], text_dim=c["text_dim"])
+    arr = np.asarray(Image.open(png_path).convert("RGB")).reshape(-1, 3)
+    total = meta["total_parameters"]
+    hi = arr[:total, 0].astype(np.uint16)
+    lo = arr[:total, 1].astype(np.uint16)
+    flat = ((hi << 8) | lo).astype(np.uint16).view(np.float16)
+    sd = dict(model.named_parameters())
+    off = 0
+    with torch.no_grad():
+        for m in meta["params"]:
+            n = m["numel"]
+            chunk = flat[off:off + n].astype(np.float16)
+            sd[m["name"]].copy_(torch.from_numpy(chunk.copy()).view(*m["shape"]).to(torch.float32))
+            off += n
+    return model.to(device).eval()
+
+def verify(weights, png_path, cfg_path, config="config.json"):
+    model = load_model_png(png_path, cfg_path)
+    ref = load_file(weights)
+    worst, name = 0.0, ""
+    for k, p in model.named_parameters():
+        d = (p.detach() - ref[k].float()).abs().max().item()
+        if d > worst:
+            worst, name = d, k
+    rel = worst / max(1e-12, ref[name].float().abs().max().item())
+    print(f"[png] round-trip max abs err {worst:.3e} on {name} (relative {rel:.2e})")
+    return worst
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--weights", default="model.safetensors")
+    ap.add_argument("--png", default="model.png")
+    ap.add_argument("--manifest", default="model_png.json")
+    ap.add_argument("--config", default="config.json")
+    args = ap.parse_args()
+    encode(args.weights, args.png, args.manifest, args.config)
+    verify(args.weights, args.png, args.manifest, args.config)
+
+if __name__ == "__main__":
+    main()

sample.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import argparse
+import base64
+import json
+import os
+
+import numpy as np
+import torch
+from safetensors.torch import load_file
+from transformers import CLIPTextModel, CLIPTokenizer
+
+from voxel_dit import VoxelDiT
+
+MAX_TOKENS = 40
+
+@torch.no_grad()
+def sample(model, seq, pool, null_seq, null_pool, steps, cfg, dev):
+    B = seq.shape[0]
+    x = torch.randn(B, 1, 32, 32, 32, device=dev)
+    ns = null_seq.expand(B, -1, -1)
+    npool = null_pool.expand(B, -1)
+    dt = 1.0 / steps
+    for i in range(steps):
+        t = torch.full((B,), i * dt, device=dev)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            vc = model(x, t, seq, pool)
+            vu = model(x, t, ns, npool)
+        x = x + (vu + cfg * (vc - vu)).float() * dt
+    return x
+
+@torch.no_grad()
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("prompts", nargs="+")
+    ap.add_argument("--out", default="samples.json")
+    ap.add_argument("--cfg", type=float, default=5.0)
+    ap.add_argument("--steps", type=int, default=50)
+    ap.add_argument("--threshold", type=float, default=0.0)
+    ap.add_argument("--device", default="cuda")
+    ap.add_argument("--weights", default="/root/runs/voxel_v1/model.safetensors")
+    ap.add_argument("--clip", default="openai/clip-vit-base-patch32")
+    args = ap.parse_args()
+    dev = args.device
+
+    model = VoxelDiT().to(dev).eval()
+    model.load_state_dict(load_file(args.weights))
+
+    tok = CLIPTokenizer.from_pretrained(args.clip)
+    txt = CLIPTextModel.from_pretrained(args.clip).to(dev).eval()
+
+    def enc(strings):
+        t = tok(strings, padding="max_length", max_length=MAX_TOKENS,
+                truncation=True, return_tensors="pt").to(dev)
+        o = txt(**t)
+        return o.last_hidden_state.float(), o.pooler_output.float()
+
+    seq, pool = enc(args.prompts)
+    null_seq, null_pool = enc([""])
+    x = sample(model, seq, pool, null_seq, null_pool, args.steps, args.cfg, dev)
+
+    grids = (x[:, 0] > args.threshold).to(torch.uint8).cpu().numpy()
+    out = []
+    for p, g in zip(args.prompts, grids):
+        occ = float(g.mean()) * 100
+        out.append({"prompt": p, "occS": round(occ, 1),
+                    "s": base64.b64encode(np.packbits(g.reshape(-1)).tobytes()).decode()})
+        print(f"  {occ:5.1f}%  {p}")
+    json.dump(out, open(args.out, "w"))
+    print(f"[sample] {len(out)} grids -> {args.out} (cfg {args.cfg}, {args.steps} steps)")
+
+if __name__ == "__main__":
+    main()

train.py

@@ -0,0 +1,186 @@
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import math
+import os
+import time
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from safetensors.torch import save_file
+
+from voxel_dit import VoxelDiT
+
+def cosine_lr(step, total, base, floor, warmup):
+    if step < warmup:
+        return base * (step + 1) / max(1, warmup)
+    t = (step - warmup) / max(1, total - warmup)
+    return floor + 0.5 * (base - floor) * (1 + math.cos(math.pi * t))
+
+def load_split(data, filtered, n_val, seed=0):
+    meta = json.load(open(f"{data}/meta.json"))
+    keep = np.array(meta["keep"], bool)
+    idx = np.nonzero(keep)[0] if filtered else np.arange(len(keep))
+    perm = np.random.RandomState(seed).permutation(len(idx))
+    idx = idx[perm]
+    val, tr = idx[:n_val], idx[n_val:]
+    packed = np.load(f"{data}/voxels_packed.npy")
+    seq = np.load(f"{data}/text_seq.npy")
+    pool = np.load(f"{data}/text_pool.npy")
+    pack = lambda s: (np.unpackbits(packed[s], axis=1), seq[s], pool[s])
+    return pack(tr), pack(val), len(keep)
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", default="/root/data")
+    ap.add_argument("--out", default="/root/runs/voxel_v1")
+    ap.add_argument("--steps", type=int, default=90000)
+    ap.add_argument("--batch-size", type=int, default=256)
+    ap.add_argument("--lr", type=float, default=2e-4)
+    ap.add_argument("--min-lr", type=float, default=1e-6)
+    ap.add_argument("--warmup", type=int, default=500)
+    ap.add_argument("--cfg-dropout", type=float, default=0.1)
+    ap.add_argument("--ema", type=float, default=0.9999)
+    ap.add_argument("--filtered", type=int, default=1)
+    ap.add_argument("--val-size", type=int, default=1024)
+    ap.add_argument("--val-every", type=int, default=2000)
+    ap.add_argument("--log-every", type=int, default=200)
+    ap.add_argument("--save-every", type=int, default=10000)
+    ap.add_argument("--resume", default="")
+    ap.add_argument("--seed", type=int, default=0)
+    args = ap.parse_args()
+
+    dev = "cuda"
+    os.makedirs(args.out, exist_ok=True)
+    torch.manual_seed(args.seed)
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+    (vox, seq, pool), (vvox, vseq, vpool), total_meshes = load_split(
+        args.data, args.filtered, args.val_size)
+    M = vox.shape[0]
+    vox_t = torch.from_numpy(vox).to(dev)
+    seq_t = torch.from_numpy(seq).to(dev)
+    pool_t = torch.from_numpy(pool).to(dev)
+    null_seq = torch.from_numpy(np.load(f"{args.data}/null_seq.npy")).to(dev).float()
+    null_pool = torch.from_numpy(np.load(f"{args.data}/null_pool.npy")).to(dev).float()
+
+    vg = torch.Generator(device=dev).manual_seed(1234)
+    vx1 = torch.from_numpy(vvox).to(dev).view(-1, 1, 32, 32, 32).float() * 2 - 1
+    vs = torch.from_numpy(vseq).to(dev).float()
+    vp = torch.from_numpy(vpool).to(dev).float()
+    vx0 = torch.randn(vx1.shape, device=dev, generator=vg)
+    vt = torch.sigmoid(torch.randn(vx1.shape[0], device=dev, generator=vg))
+
+    @torch.no_grad()
+    def val_loss():
+        ema.eval()
+        tot, n = 0.0, 0
+        for j in range(0, vx1.shape[0], args.batch_size):
+            sl = slice(j, j + args.batch_size)
+            m = vx1[sl].shape[0]
+            tb = vt[sl].view(-1, 1, 1, 1, 1)
+            xt = (1 - tb) * vx0[sl] + tb * vx1[sl]
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                v = ema(xt, vt[sl], vs[sl], vp[sl])
+            tot += F.mse_loss(v.float(), vx1[sl] - vx0[sl]).item() * m
+            n += m
+        return tot / n
+
+    model = VoxelDiT().to(dev)
+    ema = copy.deepcopy(model).eval()
+    for p in ema.parameters():
+        p.requires_grad_(False)
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.99), weight_decay=0.0)
+
+    start = 0
+    if args.resume and os.path.exists(args.resume):
+        ck = torch.load(args.resume, map_location=dev)
+        model.load_state_dict(ck["model"])
+        ema.load_state_dict(ck["ema"])
+        opt.load_state_dict(ck["opt"])
+        start = ck["step"] + 1
+
+    B = args.batch_size
+    gen = torch.Generator(device=dev).manual_seed(args.seed)
+    epochs = args.steps * B / M
+    print(f"[train] {M} train + {vx1.shape[0]} val of {total_meshes} captioned "
+          f"({'filtered' if args.filtered else 'unfiltered'}), "
+          f"{model.num_params()/1e6:.2f}M params", flush=True)
+    print(f"[train] {args.steps} steps x batch {B} = {epochs:.0f} epochs, "
+          f"resume from {start}", flush=True)
+
+    logf = open(f"{args.out}/log.jsonl", "a")
+    running, t0 = 0.0, time.time()
+    for step in range(start, args.steps):
+        i = torch.randint(0, M, (B,), device=dev, generator=gen)
+        x1 = vox_t[i].view(B, 1, 32, 32, 32).float() * 2 - 1
+        s = seq_t[i].float()
+        p = pool_t[i].float()
+        drop = torch.rand(B, device=dev, generator=gen) < args.cfg_dropout
+        s = torch.where(drop[:, None, None], null_seq, s)
+        p = torch.where(drop[:, None], null_pool, p)
+
+        x0 = torch.randn(x1.shape, device=dev, generator=gen)
+        u = torch.randn(B, device=dev, generator=gen)
+        t = torch.sigmoid(u)
+        tb = t.view(-1, 1, 1, 1, 1)
+        xt = (1 - tb) * x0 + tb * x1
+        target = x1 - x0
+
+        lr = cosine_lr(step, args.steps, args.lr, args.min_lr, args.warmup)
+        for g in opt.param_groups:
+            g["lr"] = lr
+
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            v = model(xt, t, s, p)
+            loss = F.mse_loss(v.float(), target)
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step()
+        with torch.no_grad():
+            d = 1 - args.ema
+            for pe, pm in zip(ema.parameters(), model.parameters()):
+                pe.mul_(args.ema).add_(pm.detach(), alpha=d)
+
+        running += loss.item()
+        if (step + 1) % args.log_every == 0:
+            avg = running / args.log_every
+            el = time.time() - t0
+            sps = args.log_every / el
+            eta = (args.steps - step - 1) / sps / 3600
+            print(f"[s{step+1:06d}] loss={avg:.5f} lr={lr:.2e} gnorm={gn:.2f} "
+                  f"{sps:.2f} steps/s eta={eta:.2f}h", flush=True)
+            logf.write(json.dumps({"step": step + 1, "loss": avg, "lr": lr,
+                                   "gnorm": float(gn), "steps_per_s": sps}) + "\n")
+            logf.flush()
+            running, t0 = 0.0, time.time()
+
+        if (step + 1) % args.val_every == 0 or step + 1 == args.steps:
+            vl = val_loss()
+            print(f"[s{step+1:06d}] val_loss={vl:.5f} (ema, {args.val_size} held out)", flush=True)
+            logf.write(json.dumps({"step": step + 1, "val_loss": vl}) + "\n")
+            logf.flush()
+            t0 = time.time()
+
+        if (step + 1) % args.save_every == 0 or step + 1 == args.steps:
+            torch.save({"model": model.state_dict(), "ema": ema.state_dict(),
+                        "opt": opt.state_dict(), "step": step},
+                       f"{args.out}/ckpt.pt")
+            save_file({k: v.contiguous() for k, v in ema.state_dict().items()},
+                      f"{args.out}/model.safetensors")
+            print(f"[s{step+1:06d}] saved", flush=True)
+
+    json.dump({"dit": {"vox_size": 32, "patch": 4, "dim": 384, "depth": 12,
+                       "heads": 6, "text_dim": 512},
+               "steps": args.steps, "batch_size": args.batch_size,
+               "meshes": M, "filtered": bool(args.filtered)},
+              open(f"{args.out}/config.json", "w"), indent=2)
+    print("TRAINDONE", flush=True)
+
+if __name__ == "__main__":
+    main()

voxel_dit.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+def timestep_embedding(t, dim, max_period=10000):
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(half, device=t.device) / half)
+    args = t[:, None].float() * freqs[None]
+    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
+    return emb
+
+def sincos_3d(dim, grid_size):
+    assert dim % 6 == 0
+    g = np.arange(grid_size, dtype=np.float32)
+    gz, gy, gx = np.meshgrid(g, g, g, indexing="ij")
+    d6 = dim // 6
+    omega = 1.0 / (10000 ** (np.arange(d6, dtype=np.float32) / d6))
+    def emb1(p):
+        out = p.reshape(-1)[:, None] * omega[None]
+        return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+    pe = np.concatenate([emb1(gz), emb1(gy), emb1(gx)], axis=1)
+    return torch.from_numpy(pe).float()
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.q = nn.Linear(dim, dim)
+        self.kv = nn.Linear(dim, dim * 2)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x, ctx=None):
+        ctx = x if ctx is None else ctx
+        B, N, C = x.shape
+        M = ctx.shape[1]
+        h = self.heads
+        q = self.q(x).reshape(B, N, h, C // h).transpose(1, 2)
+        kv = self.kv(ctx).reshape(B, M, 2, h, C // h).permute(2, 0, 3, 1, 4)
+        k, v = kv[0], kv[1]
+        o = F.scaled_dot_product_attention(q, k, v)
+        o = o.transpose(1, 2).reshape(B, N, C)
+        return self.proj(o)
+
+class Block(nn.Module):
+    def __init__(self, dim, heads, mlp_ratio=4.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(dim, heads)
+        self.norm_ca = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.cross = Attention(dim, heads)
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        hidden = int(dim * mlp_ratio)
+        self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(approximate="tanh"),
+                                 nn.Linear(hidden, dim))
+        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
+        self.cross_gate = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x, c, text):
+        shift1, scale1, gate1, shift2, scale2, gate2 = self.ada(c).chunk(6, dim=1)
+        x = x + gate1.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift1, scale1))
+        x = x + self.cross_gate * self.cross(self.norm_ca(x), text)
+        x = x + gate2.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift2, scale2))
+        return x
+
+class VoxelDiT(nn.Module):
+    def __init__(self, vox_ch=1, vox_size=32, patch=4, dim=384, depth=12,
+                 heads=6, text_dim=512, mlp_ratio=4.0):
+        super().__init__()
+        self.vox_ch = vox_ch
+        self.vox_size = vox_size
+        self.patch = patch
+        self.grid = vox_size // patch
+        self.patch_dim = vox_ch * patch * patch * patch
+        self.x_embed = nn.Linear(self.patch_dim, dim)
+        self.register_buffer("pos", sincos_3d(dim, self.grid).unsqueeze(0))
+        self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.text_proj = nn.Linear(text_dim, dim)
+        self.text_pool = nn.Linear(text_dim, dim)
+        self.blocks = nn.ModuleList([Block(dim, heads, mlp_ratio) for _ in range(depth)])
+        self.norm_out = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ada_out = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim))
+        self.head = nn.Linear(dim, self.patch_dim)
+        self.dim = dim
+        self._init()
+
+    def _init(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+        for b in self.blocks:
+            nn.init.zeros_(b.ada[-1].weight); nn.init.zeros_(b.ada[-1].bias)
+        nn.init.zeros_(self.ada_out[-1].weight); nn.init.zeros_(self.ada_out[-1].bias)
+        nn.init.zeros_(self.head.weight); nn.init.zeros_(self.head.bias)
+
+    def patchify(self, x):
+        B, C, D, H, W = x.shape
+        p = self.patch
+        x = x.reshape(B, C, D // p, p, H // p, p, W // p, p)
+        x = x.permute(0, 2, 4, 6, 1, 3, 5, 7)
+        return x.reshape(B, (D // p) * (H // p) * (W // p), C * p * p * p)
+
+    def unpatchify(self, x):
+        B, N, _ = x.shape
+        p = self.patch
+        g = self.grid
+        C = self.vox_ch
+        x = x.reshape(B, g, g, g, C, p, p, p)
+        x = x.permute(0, 4, 1, 5, 2, 6, 3, 7)
+        return x.reshape(B, C, g * p, g * p, g * p)
+
+    def forward(self, x, t, text_seq, text_pool):
+        x = self.x_embed(self.patchify(x)) + self.pos
+        c = self.t_mlp(timestep_embedding(t, self.dim)) + self.text_pool(text_pool)
+        text = self.text_proj(text_seq)
+        for blk in self.blocks:
+            x = blk(x, c, text)
+        shift, scale = self.ada_out(c).chunk(2, dim=1)
+        x = modulate(self.norm_out(x), shift, scale)
+        x = self.head(x)
+        return self.unpatchify(x)
+
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())