diff --git a/README.md b/README.md index 590328b..9e18b20 100644 --- a/README.md +++ b/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 diff --git a/pixal3d_mlx/decoders.py b/pixal3d_mlx/decoders.py new file mode 100644 index 0000000..a11acd3 --- /dev/null +++ b/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...<...>; 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 diff --git a/tests/test_decoders.py b/tests/test_decoders.py new file mode 100644 index 0000000..1d06d30 --- /dev/null +++ b/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)