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.
This commit is contained in:
John
parent
8faf6d592a
commit
6ac586b947
@ -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
|
||||
|
||||
206
pixal3d_mlx/ss_dec.py
Normal file
206
pixal3d_mlx/ss_dec.py
Normal file
@ -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
|
||||
Loading…
Reference in New Issue
Block a user