From 6ac586b9470ca17a3fd5d2d5a792b278d19049c2 Mon Sep 17 00:00:00 2001 From: John Date: Sun, 2 Aug 2026 13:27:24 +1000 Subject: [PATCH] SparseStructureDecoder verified at correlation 1.00000000 74/74 params, max abs diff 4.6e-4 on values around -163 (~3e-6 relative), 60ms for a 16^3 latent -> 64^3 occupancy grid. Fully dense, so no sparse ops involved - but MLX's Conv3d is channels-LAST where torch is NCDHW, so tensors are carried channels-last throughout and transposed only at the boundaries. The converter already emits [O,kz,ky,kx,I] to match. pixel_shuffle_3d had to be rewritten for that layout: the (H,s)(W,s)(D,s) interleave order is what matters and getting it wrong scrambles the grid while preserving its shape. This completes the structure stage end to end: image -> ss_flow -> latent -> ss_dec. --- README.md | 9 +- pixal3d_mlx/ss_dec.py | 206 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 212 insertions(+), 3 deletions(-) create mode 100644 pixal3d_mlx/ss_dec.py diff --git a/README.md b/README.md index 7ccbb4b..590328b 100644 --- a/README.md +++ b/README.md @@ -75,9 +75,12 @@ Weights: 24.04 GB across 19 files (1.3B DiTs at 512/1024 + shape/tex decoders). - [x] **All four flow models verified against upstream at correlation 1.00000000** — `ss_flow` (max diff 1.2e-5) and `slat_flow` (9.3e-6), 700/700 params each, on the real 1.3B checkpoints -- [ ] `SparseConvNeXtBlock3d`, `SparseResBlockC2S3d`, `SparseSpatial2Channel` (decoders) -- [ ] Model graphs -- [ ] End-to-end +- [x] `SparseConvNeXtBlock3d`, `SparseResBlockC2S3d`, `spatial2channel`/`channel2spatial` + — 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 ## Numerical verification diff --git a/pixal3d_mlx/ss_dec.py b/pixal3d_mlx/ss_dec.py new file mode 100644 index 0000000..94aa16f --- /dev/null +++ b/pixal3d_mlx/ss_dec.py @@ -0,0 +1,206 @@ +"""Pixal3D SparseStructureDecoder in MLX — latent -> occupancy grid. + +Fully dense: 3D convs on a small voxel grid, no sparse ops at all. Pairs with `ss_flow` +to complete the structure stage (image -> latent -> occupancy), which is what decides +which voxels exist before the SLAT stage runs. + +Config (from the sibling JSON): + out_channels 1, latent_channels 8, channels [512, 128, 32], + num_res_blocks 2, num_res_blocks_middle 2, norm_type "layer" + +MLX's Conv3d is channels-LAST (NDHWC) where torch is NCDHW, so tensors are carried in +channels-last form throughout and only transposed at the boundaries. The weight +converter already emits `[O, kz, ky, kx, I]` to match. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import List, Optional + +import mlx.core as mx +import mlx.nn as nn + + +def _cln(x: mx.array, weight: mx.array, bias: mx.array, eps: float = 1e-5) -> mx.array: + """ChannelLayerNorm32 — LayerNorm over channels, computed in fp32. + + Upstream permutes channels to last, applies LayerNorm, permutes back. Data here is + already channels-last, so it is a plain LayerNorm over the final axis. + """ + dt = x.dtype + f = x.astype(mx.float32) + f = (f - mx.mean(f, -1, keepdims=True)) * mx.rsqrt(mx.var(f, -1, keepdims=True) + eps) + return (f * weight + bias).astype(dt) + + +def pixel_shuffle_3d(x: mx.array, scale: int) -> mx.array: + """Channels-last 3D pixel shuffle: [B,H,W,D,C*s^3] -> [B,H*s,W*s,D*s,C]. + + Upstream works in NCDHW and permutes (0,1,5,2,6,3,7,4) after splitting the channel + axis into (C_, s, s, s). Here the channel axis is last, so the split factors sit + next to it and the interleave order must still be (H,s)(W,s)(D,s) — getting that + order wrong scrambles the grid while preserving its shape. + """ + b, h, w, d, c = x.shape + c_ = c // scale**3 + x = x.reshape(b, h, w, d, c_, scale, scale, scale) + x = x.transpose(0, 1, 5, 2, 6, 3, 7, 4) # b, h, s, w, s, d, s, c_ + return x.reshape(b, h * scale, w * scale, d * scale, c_) + + +class ResBlock3d(nn.Module): + def __init__(self, channels: int, out_channels: Optional[int] = None): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.norm1_w = mx.ones((channels,)) + self.norm1_b = mx.zeros((channels,)) + self.norm2_w = mx.ones((self.out_channels,)) + self.norm2_b = mx.zeros((self.out_channels,)) + self.conv1 = nn.Conv3d(channels, self.out_channels, 3, padding=1) + self.conv2 = nn.Conv3d(self.out_channels, self.out_channels, 3, padding=1) + self.skip_connection = ( + nn.Conv3d(channels, self.out_channels, 1) + if channels != self.out_channels + else None + ) + + def __call__(self, x: mx.array) -> mx.array: + h = _cln(x, self.norm1_w, self.norm1_b) + h = self.conv1(nn.silu(h)) + h = _cln(h, self.norm2_w, self.norm2_b) + h = self.conv2(nn.silu(h)) + return h + (self.skip_connection(x) if self.skip_connection else x) + + +class UpsampleBlock3d(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.conv = nn.Conv3d(in_channels, out_channels * 8, 3, padding=1) + + def __call__(self, x: mx.array) -> mx.array: + return pixel_shuffle_3d(self.conv(x), 2) + + +class SparseStructureDecoder(nn.Module): + def __init__( + self, + out_channels: int = 1, + latent_channels: int = 8, + num_res_blocks: int = 2, + channels: List[int] = (512, 128, 32), + num_res_blocks_middle: int = 2, + **_ignored, + ): + super().__init__() + channels = list(channels) + self.channels = channels + self.input_layer = nn.Conv3d(latent_channels, channels[0], 3, padding=1) + self.middle_block = [ + ResBlock3d(channels[0], channels[0]) for _ in range(num_res_blocks_middle) + ] + blocks = [] + for i, ch in enumerate(channels): + blocks += [ResBlock3d(ch, ch) for _ in range(num_res_blocks)] + if i < len(channels) - 1: + blocks.append(UpsampleBlock3d(ch, channels[i + 1])) + self.blocks = blocks + self.out_norm_w = mx.ones((channels[-1],)) + self.out_norm_b = mx.zeros((channels[-1],)) + self.out_conv = nn.Conv3d(channels[-1], out_channels, 3, padding=1) + + def __call__(self, z: mx.array) -> mx.array: + """z: [B, C, D, H, W] (torch order, for a familiar interface).""" + x = z.transpose(0, 2, 3, 4, 1) # -> channels-last + x = self.input_layer(x) + for b in self.middle_block: + x = b(x) + for b in self.blocks: + x = b(x) + x = _cln(x, self.out_norm_w, self.out_norm_b) + x = self.out_conv(nn.silu(x)) + return x.transpose(0, 4, 1, 2, 3) # back to [B, C, D, H, W] + + +# ------------------------------------------------------------------ loading + +def _remap(k: str) -> str: + """Upstream nests norms/convs inside Sequential and nn.Module attrs.""" + k = k.replace("out_layer.0.weight", "out_norm_w").replace("out_layer.0.bias", "out_norm_b") + k = k.replace("out_layer.2.", "out_conv.") + k = k.replace(".norm1.weight", ".norm1_w").replace(".norm1.bias", ".norm1_b") + k = k.replace(".norm2.weight", ".norm2_w").replace(".norm2.bias", ".norm2_b") + return k + + +def load(weights_path: str | Path, config_path: str | Path | None = None): + wp = Path(weights_path) + cp = Path(config_path) if config_path else wp.with_suffix(".json") + cfg = json.loads(cp.read_text()) + args = dict(cfg.get("args", {})) + args.pop("use_fp16", None) + model = SparseStructureDecoder(**args) + + w = mx.load(str(wp)) + flat = dict(_flatten(model.parameters())) + mapped, unmapped = {}, [] + for k, v in w.items(): + m = _remap(k) + if m in flat: + if flat[m].shape != v.shape: + raise ValueError(f"shape mismatch {k} -> {m}: {flat[m].shape} vs {v.shape}") + 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": cfg.get("name"), + "loaded": len(mapped), + "params": len(flat), + "missing": missing, + "unmapped": unmapped, + } + + +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