Sparse VAE decoders: shape_dec and tex_dec

All seven Pixal3D models now load and run. shape_dec 292/292, tex_dec 284/284, both
with zero missing/unmapped/mismatched keys - the 8-param difference between them is
exactly the four to_subdiv layers, since tex_dec has pred_subdiv=False.

Behaviour is right: 12 latent voxels grow SELECTIVELY through four stages
(12 -> 91 -> 193 -> 358 -> 1150) rather than x8 each time, which would have reached
49,152. Scale lands at 1/16. shape_dec emits the hardcoded 7 channels and its vertex
head produces offsets inside the [-0.5, 1.5] band its sigmoid+voxel_margin allows.
tex_dec, guided by shape_dec's masks, reproduces exactly the same voxel count.

VERIFICATION CAVEAT, recorded in the README: these two are the only models using sparse
conv, so upstream cannot run here and there is no numerical oracle. Unlike the five
models verified at correlation 1.0, these are checked structurally and behaviourally
only. Weaker evidence, and labelled as such rather than presented alongside the
verified results.

README.md

@@ -79,10 +79,28 @@ Weights: 24.04 GB across 19 files (1.3B DiTs at 512/1024 + shape/tex decoders).
       — all in the shared core, round-trip and selective-growth tested
 - [x] **`ss_dec` verified at correlation 1.00000000** (74/74), completing the whole
       structure stage: image -> ss_flow -> latent -> ss_dec -> 64^3 occupancy grid
-- [ ] `shape_dec` / `tex_dec` model graphs (the two sparse decoders)
-- [ ] End-to-end pipeline wiring
+- [x] `shape_dec` / `tex_dec` — both load complete (292/292, 284/284) and run.
+      **Behaviourally** checked only; see the verification note below
+- [ ] End-to-end pipeline wiring (image encoder -> flows -> decoders -> mesh export)
 
 
+## Model status
+
+| model | params | verification |
+|---|---|---|
+| `ss_flow` | 700/700 | **corr 1.00000000** vs upstream |
+| `slat_flow` x3 | 700/700 | **corr 1.00000000** vs upstream |
+| `ss_dec` | 74/74 | **corr 1.00000000** vs upstream |
+| `shape_dec` | 292/292 | behavioural only — no oracle possible |
+| `tex_dec` | 284/284 | behavioural only — no oracle possible |
+
+The split is not arbitrary: the first five contain no sparse convolution, so upstream
+runs on CPU torch and can be diffed directly. The two decoders do use it, spconv has no
+Metal build, and so there is nothing to diff against. Their blocks are individually
+tested, and the assembled graphs are checked for complete weight mapping, selective
+growth, correct scale and a vertex head inside its valid band — but that is weaker
+evidence than a correlation and should be read that way.
+
 ## Numerical verification
 
 The flow models contain no sparse convolution, which means **upstream runs on CPU torch

pixal3d_mlx/decoders.py

@@ -0,0 +1,244 @@
+"""Pixal3D sparse VAE decoders in MLX — `shape_dec` and `tex_dec`.
+
+Both are the same U-Net: stages of `SparseConvNeXtBlock3d`, each stage ending in a
+`SparseResBlockC2S3d` that upsamples via channel->spatial. The occupied set grows
+selectively at every stage, driven by a predicted subdivision mask, which is how the
+decoder reaches 256^3-equivalent detail without ever materialising a dense grid.
+
+    model_channels [1024, 512, 256, 128, 64]
+    num_blocks     [4, 16, 8, 4, 0]        <- last stage is empty and has no up-block
+    block_type     SparseConvNeXtBlock3d
+    up_block_type  SparseResBlockC2S3d
+
+`FlexiDualGridVaeDecoder` (shape_dec) subclasses it and hardcodes **7** output channels
+regardless of config — 3 for vertex offsets (squashed through a sigmoid widened by
+`voxel_margin`) plus 4 more — so its `out_channels` is not a config field at all.
+`SparseUnetVaeDecoder` (tex_dec) uses the config's `out_channels` (6) and sets
+`pred_subdiv=False`, meaning it needs subdivision masks supplied from outside rather
+than predicting its own.
+
+NOTE ON VERIFICATION: unlike the flow models and ss_dec, these use sparse convolution,
+so upstream cannot run here (no spconv on Metal) and there is no numerical oracle. They
+are built from blocks that ARE individually tested, but the assembled graphs are checked
+only for shape, finiteness and plausible behaviour. Treat them as less certain than
+anything reporting a correlation.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from trellis_sparse_mlx import (
+    SparseConvNeXtBlock3d,
+    SparseLinear,
+    SparseResBlockC2S3d,
+    SparseTensor,
+)
+
+_BLOCKS = {
+    "SparseConvNeXtBlock3d": SparseConvNeXtBlock3d,
+    "SparseResBlockC2S3d": SparseResBlockC2S3d,
+}
+
+
+class SparseUnetVaeDecoder(nn.Module):
+    def __init__(
+        self,
+        out_channels: int,
+        model_channels: List[int],
+        latent_channels: int,
+        num_blocks: List[int],
+        block_type: List[str],
+        up_block_type: List[str],
+        block_args: Optional[List[Dict[str, Any]]] = None,
+        pred_subdiv: bool = True,
+        **_ignored,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.model_channels = list(model_channels)
+        self.num_blocks = list(num_blocks)
+        self.pred_subdiv = pred_subdiv
+        block_args = block_args or [{} for _ in num_blocks]
+
+        self.from_latent = SparseLinear(latent_channels, model_channels[0])
+        self.output_layer = SparseLinear(model_channels[-1], out_channels)
+
+        stages = []
+        for i in range(len(num_blocks)):
+            stage = [
+                _BLOCKS[block_type[i]](model_channels[i], **block_args[i])
+                for _ in range(num_blocks[i])
+            ]
+            if i < len(num_blocks) - 1:
+                stage.append(
+                    _BLOCKS[up_block_type[i]](
+                        model_channels[i],
+                        model_channels[i + 1],
+                        pred_subdiv=pred_subdiv,
+                    )
+                )
+            stages.append(stage)
+        self.blocks = stages
+
+    def __call__(
+        self,
+        x: SparseTensor,
+        guide_subs: Optional[List[SparseTensor]] = None,
+        return_subs: bool = False,
+    ):
+        h = self.from_latent(x)
+        subs = []
+        for i, stage in enumerate(self.blocks):
+            for j, blk in enumerate(stage):
+                is_up = i < len(self.blocks) - 1 and j == len(stage) - 1
+                if is_up:
+                    if self.pred_subdiv:
+                        h, sub = blk(h)
+                        subs.append(sub)
+                    else:
+                        h = blk(h, subdiv=guide_subs[i] if guide_subs else None)
+                else:
+                    h = blk(h)
+
+        # Same parameterless final LayerNorm as every other model in this family.
+        f = h.feats
+        f = (f - mx.mean(f, -1, keepdims=True)) * mx.rsqrt(
+            mx.var(f, -1, keepdims=True) + 1e-5
+        )
+        out = self.output_layer(h.replace(f))
+        return (out, subs) if return_subs else out
+
+
+class FlexiDualGridVaeDecoder(SparseUnetVaeDecoder):
+    """shape_dec. Emits 7 channels; the first 3 are vertex offsets within the voxel."""
+
+    def __init__(
+        self,
+        resolution: int,
+        model_channels: List[int],
+        latent_channels: int,
+        num_blocks: List[int],
+        block_type: List[str],
+        up_block_type: List[str],
+        block_args: Optional[List[Dict[str, Any]]] = None,
+        voxel_margin: float = 0.5,
+        **_ignored,
+    ):
+        # out_channels is hardcoded to 7 upstream, not read from config
+        super().__init__(
+            7,
+            model_channels,
+            latent_channels,
+            num_blocks,
+            block_type,
+            up_block_type,
+            block_args,
+            pred_subdiv=True,
+        )
+        self.resolution = resolution
+        self.voxel_margin = voxel_margin
+
+    def decode_vertices(self, h: SparseTensor) -> mx.array:
+        """Channels 0:3 -> vertex offsets, sigmoid widened by `voxel_margin`.
+
+        `(1 + 2m)*sigmoid(v) - m` maps to [-m, 1+m], letting a vertex sit slightly
+        outside its own voxel — that overshoot is what lets the dual grid represent
+        surfaces that do not pass exactly through cell centres.
+        """
+        m = self.voxel_margin
+        return (1 + 2 * m) * mx.sigmoid(h.feats[..., 0:3]) - m
+
+
+# ------------------------------------------------------------------ loading
+
+def _remap(k: str) -> str:
+    """Upstream: blocks.<stage>.<idx>.<...>; SparseLinear wraps an nn.Linear."""
+    for name in ("from_latent", "output_layer", "to_subdiv"):
+        k = k.replace(f"{name}.weight", f"{name}.linear.weight")
+        k = k.replace(f"{name}.bias", f"{name}.linear.bias")
+    k = k.replace(".mlp.0.", ".mlp_0.").replace(".mlp.2.", ".mlp_2.")
+    k = k.replace(".conv1.conv.", ".conv1.").replace(".conv2.conv.", ".conv2.")
+    k = k.replace(".conv.conv.", ".conv.")
+    k = k.replace(".norm.weight", ".norm.weight").replace(".norm1.weight", ".norm1.weight")
+    return k
+
+
+def load(weights_path: str | Path, config_path: str | Path):
+    cfg = json.loads(Path(config_path).read_text())
+    name, args = cfg["name"], dict(cfg.get("args", {}))
+    args.pop("use_fp16", None)
+    cls = {
+        "FlexiDualGridVaeDecoder": FlexiDualGridVaeDecoder,
+        "SparseUnetVaeDecoder": SparseUnetVaeDecoder,
+    }[name]
+    model = cls(**args)
+
+    w = mx.load(str(weights_path))
+    flat = dict(_flatten(model.parameters()))
+    mapped, unmapped, mismatched = {}, [], []
+    for k, v in w.items():
+        m = _remap(k)
+        if m in flat:
+            if flat[m].shape != v.shape:
+                mismatched.append(f"{k} -> {m}: want {flat[m].shape} got {v.shape}")
+                continue
+            mapped[m] = v
+        else:
+            unmapped.append(f"{k} -> {m} {tuple(v.shape)}")
+    missing = [k for k in flat if k not in mapped]
+    if mapped:
+        model.update(_unflatten(mapped))
+    return model, {
+        "config": name,
+        "loaded": len(mapped),
+        "params": len(flat),
+        "missing": missing,
+        "unmapped": unmapped,
+        "mismatched": mismatched,
+    }
+
+
+def _flatten(tree, prefix=""):
+    if isinstance(tree, dict):
+        for k, v in tree.items():
+            yield from _flatten(v, f"{prefix}{k}.")
+    elif isinstance(tree, list):
+        for i, v in enumerate(tree):
+            yield from _flatten(v, f"{prefix}{i}.")
+    elif isinstance(tree, mx.array):
+        yield prefix[:-1], tree
+
+
+def _unflatten(flat: dict):
+    root: dict = {}
+    for key, val in flat.items():
+        parts = key.split(".")
+        node = root
+        for i, p in enumerate(parts[:-1]):
+            nxt = parts[i + 1]
+            default = [] if nxt.isdigit() else {}
+            if isinstance(node, list):
+                idx = int(p)
+                while len(node) <= idx:
+                    node.append({})
+                if isinstance(default, list) and not isinstance(node[idx], list):
+                    node[idx] = default
+                node = node[idx]
+            else:
+                if p not in node or not isinstance(node[p], (dict, list)):
+                    node[p] = default
+                node = node[p]
+        if isinstance(node, list):
+            idx = int(parts[-1])
+            while len(node) <= idx:
+                node.append(None)
+            node[idx] = val
+        else:
+            node[parts[-1]] = val
+    return root

tests/test_decoders.py

@@ -0,0 +1,110 @@
+"""Behavioural tests for the two sparse decoders.
+
+These are the only Pixal3D models using sparse convolution, so upstream cannot run here
+(spconv has no Metal build) and there is NO numerical oracle — unlike the flow models and
+ss_dec, which are all verified at correlation 1.0 against upstream on CPU torch.
+
+So these check structure and behaviour rather than values: that weights map completely,
+that the occupied set grows SELECTIVELY rather than x8, that scale tracks the four
+upsamples, and that the vertex head lands inside the band its sigmoid+margin allows.
+Weaker than a correlation, and stated as such.
+"""
+
+import sys
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from pixal3d_mlx.decoders import load  # noqa: E402
+from trellis_sparse_mlx import SparseTensor  # noqa: E402
+
+CK = Path("/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts")
+W = Path(__file__).resolve().parents[1] / "weights"
+
+
+def _latent(n=12, ch=32, res=4, seed=0):
+    rng = np.random.default_rng(seed)
+    co = np.concatenate(
+        [np.zeros((n, 1), dtype=np.int32), rng.integers(0, res, (n, 3)).astype(np.int32)], 1
+    )
+    co = np.unique(co, axis=0)
+    return SparseTensor(
+        mx.array(rng.standard_normal((len(co), ch)).astype(np.float32)), mx.array(co)
+    )
+
+
+def _have(stem):
+    return (W / f"{stem}.safetensors").exists() and (CK / f"{stem}.json").exists()
+
+
+def test_shape_dec():
+    stem = "shape_dec_next_dc_f16c32_fp16"
+    if not _have(stem):
+        print("      (skipped: weights absent)")
+        return 0.0
+    m, rep = load(W / f"{stem}.safetensors", CK / f"{stem}.json")
+    assert not rep["missing"] and not rep["unmapped"] and not rep["mismatched"], rep
+    x = _latent()
+    n_in = x.coords.shape[0]
+    out, subs = m(x, return_subs=True)
+    f = np.asarray(out.feats)
+
+    assert f.shape[1] == 7, f"FlexiDualGrid emits 7 channels, got {f.shape[1]}"
+    assert np.isfinite(f).all(), "non-finite output"
+    # four x2 upsamples -> scale 1/16
+    assert abs(out._scale[0] - 1 / 16) < 1e-9, out._scale
+    # growth must be selective, not the full x8 per stage
+    assert out.coords.shape[0] > n_in, "no growth at all"
+    assert out.coords.shape[0] < n_in * 8**4, "grew by the full x8 every stage"
+    assert len(subs) == 4, f"expected 4 subdivision masks, got {len(subs)}"
+
+    v = np.asarray(m.decode_vertices(out))
+    m_ = m.voxel_margin
+    assert v.min() >= -m_ - 1e-4 and v.max() <= 1 + m_ + 1e-4, (v.min(), v.max())
+    return 0.0
+
+
+def test_tex_dec_follows_guide_masks():
+    """tex_dec has pred_subdiv=False, so its structure must match the masks it is given."""
+    s_stem, t_stem = "shape_dec_next_dc_f16c32_fp16", "tex_dec_next_dc_f16c32_fp16"
+    if not (_have(s_stem) and _have(t_stem)):
+        print("      (skipped: weights absent)")
+        return 0.0
+    sm, _ = load(W / f"{s_stem}.safetensors", CK / f"{s_stem}.json")
+    tm, rep = load(W / f"{t_stem}.safetensors", CK / f"{t_stem}.json")
+    assert not rep["missing"] and not rep["unmapped"], rep
+
+    x = _latent()
+    shape_out, subs = sm(x, return_subs=True)
+    tex_out = tm(_latent(), guide_subs=subs)
+    f = np.asarray(tex_out.feats)
+
+    assert f.shape[1] == 6, f"tex_dec emits 6 channels, got {f.shape[1]}"
+    assert np.isfinite(f).all(), "non-finite output"
+    assert tex_out.coords.shape[0] == shape_out.coords.shape[0], (
+        "guided decoder produced a different voxel count than the masks describe"
+    )
+    return 0.0
+
+
+if __name__ == "__main__":
+    tests = [
+        ("shape_dec structure", test_shape_dec),
+        ("tex_dec follows guides", test_tex_dec_follows_guide_masks),
+    ]
+    failed = 0
+    for name, fn in tests:
+        try:
+            fn()
+            print(f"  PASS  {name}")
+        except AssertionError as e:
+            print(f"  FAIL  {name}: {e}")
+            failed += 1
+        except Exception as e:  # noqa: BLE001
+            print(f"  ERROR {name}: {type(e).__name__}: {e}")
+            failed += 1
+    print(f"\n{len(tests)-failed}/{len(tests)} passed")
+    sys.exit(1 if failed else 0)