From 8faf6d592a218f120672ae28c50e20ef72dabd95 Mon Sep 17 00:00:00 2001 From: John Date: Sun, 2 Aug 2026 12:10:28 +1000 Subject: [PATCH] All four flow models verified at correlation 1.00000000 slat_flow joins ss_flow: max abs diff 9.3e-6, 700/700 params, against upstream running on CPU torch. All three SLAT checkpoints load clean and run (img2shape 512/1024 and imgshape2tex, the last taking 64 in-channels since shape is concatenated). The SLAT flows differ from ss_flow in two ways, both handled in the shared core: tokens are a SparseTensor's voxels so attention runs per batch item, and RoPE phases are NOT shipped - positions are the input's own coordinates, so they are derived at call time by rope_phases_from_coords. tests/oracle_slat.py keeps the CPU-torch oracle harness: it patches both the dense and sparse flash-attn kernels with SDPA equivalents. Use it rather than reasoning about correctness - it has already caught three bugs a perfect 700/700 key match did not. --- README.md | 7 +- pixal3d_mlx/slat_flow.py | 192 +++++++++++++++++++++++++++++++++++++++ tests/oracle_slat.py | 108 ++++++++++++++++++++++ 3 files changed, 304 insertions(+), 3 deletions(-) create mode 100644 pixal3d_mlx/slat_flow.py create mode 100644 tests/oracle_slat.py diff --git a/README.md b/README.md index eb82651..7ccbb4b 100644 --- a/README.md +++ b/README.md @@ -72,9 +72,10 @@ Weights: 24.04 GB across 19 files (1.3B DiTs at 512/1024 + shape/tex decoders). - [x] **Weight converter** — decoders remap KRSC→`[K³,in,out]`, dtype preserved, remap verified a pure permutation. Flow models pass through untouched. - [x] DiT blocks: RoPE (stored complex phases), qk_rms_norm, AdaLN modulation, proj conditioning -- [x] **`SparseStructureFlowModel` verified against upstream: correlation 1.00000000**, - max abs diff 1.2e-5, on the real 1.3B checkpoint (700/700 params) -- [ ] `SparseConvNeXtBlock3d`, `SparseResBlockC2S3d`, `SparseSpatial2Channel` +- [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 diff --git a/pixal3d_mlx/slat_flow.py b/pixal3d_mlx/slat_flow.py new file mode 100644 index 0000000..9e85391 --- /dev/null +++ b/pixal3d_mlx/slat_flow.py @@ -0,0 +1,192 @@ +"""Pixal3D ElasticSLatFlowModel in MLX — the three SLAT flow checkpoints. + +Same DiT as `ss_flow`, with two differences: + +* Tokens are a SparseTensor's voxels, not a dense grid, so attention runs within each + batch item and `input_layer`/`out_layer` are sparse linears. +* RoPE phases are NOT shipped. `ss_flow` carries a precomputed `rope_phases` for its + fixed 16³ grid; here positions are the input's own coordinates and vary per call, so + they are derived with `rope_phases_from_coords` (verified against ss_flow's shipped + tensor to 9.6e-7). + +Three checkpoints share this class: + img2shape_512 resolution 32, in 32 -> out 32 + img2shape_1024 resolution 64, in 32 -> out 32 + imgshape2tex resolution 64, in 64 -> out 32 (shape is concatenated in) +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Optional + +import mlx.core as mx +import mlx.nn as nn + +from trellis_sparse_mlx import ( + ModulatedSparseTransformerCrossBlock, + SparseLinear, + SparseTensor, + TimestepEmbedder, + rope_phases_from_coords, +) + + +class SLatFlowModel(nn.Module): + def __init__( + self, + resolution: int = 64, + in_channels: int = 32, + out_channels: int = 32, + model_channels: int = 1536, + cond_channels: int = 1024, + num_blocks: int = 30, + num_heads: int = 12, + mlp_ratio: float = 5.3334, + share_mod: bool = True, + qk_rms_norm: bool = True, + qk_rms_norm_cross: bool = True, + image_attn_mode: str = "proj", + proj_in_channels: Optional[int] = None, + pe_mode: str = "rope", + **_ignored, + ): + super().__init__() + self.resolution = resolution + self.in_channels = in_channels + self.out_channels = out_channels + self.model_channels = model_channels + self.num_heads = num_heads + self.share_mod = share_mod + self.pe_mode = pe_mode + + self.t_embedder = TimestepEmbedder(model_channels) + if share_mod: + self.adaLN_modulation = nn.Linear(model_channels, 6 * model_channels) + self.input_layer = SparseLinear(in_channels, model_channels) + self.blocks = [ + ModulatedSparseTransformerCrossBlock( + model_channels, + cond_channels, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + share_mod=share_mod, + use_rope=(pe_mode == "rope"), + qk_rms_norm=qk_rms_norm, + qk_rms_norm_cross=qk_rms_norm_cross, + image_attn_mode=image_attn_mode, + proj_in_channels=proj_in_channels, + ) + for _ in range(num_blocks) + ] + self.out_layer = SparseLinear(model_channels, out_channels) + + def __call__(self, x: SparseTensor, t: mx.array, cond) -> SparseTensor: + h = self.input_layer(x) + + t_emb = self.t_embedder(t) + if self.share_mod: + t_emb = self.adaLN_modulation(nn.silu(t_emb)) + + phases = None + if self.pe_mode == "rope": + phases = rope_phases_from_coords( + x.coords[:, 1:], head_dim=self.model_channels // self.num_heads + ) + + for blk in self.blocks: + h = blk(h, t_emb, cond, phases=phases) + + # Same parameterless final LayerNorm as ss_flow — no params, so no checkpoint + # trace; without it the output is orders of magnitude too large. + f = h.feats + f = (f - mx.mean(f, -1, keepdims=True)) * mx.rsqrt( + mx.var(f, -1, keepdims=True) + 1e-5 + ) + return self.out_layer(h.replace(f)) + + +def _remap(k: str) -> str: + k = k.replace("t_embedder.mlp.0.", "t_embedder.mlp_0.") + k = k.replace("t_embedder.mlp.2.", "t_embedder.mlp_2.") + k = k.replace(".mlp.mlp.0.", ".mlp_0.") + k = k.replace(".mlp.mlp.2.", ".mlp_2.") + k = k.replace("adaLN_modulation.1.", "adaLN_modulation.") + # our SparseLinear wraps an nn.Linear + for name in ("input_layer", "out_layer"): + k = k.replace(f"{name}.weight", f"{name}.linear.weight") + k = k.replace(f"{name}.bias", f"{name}.linear.bias") + 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("dtype", None) + args.pop("initialization", None) + + model = SLatFlowModel(**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}") + 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 diff --git a/tests/oracle_slat.py b/tests/oracle_slat.py new file mode 100644 index 0000000..6a2ba29 --- /dev/null +++ b/tests/oracle_slat.py @@ -0,0 +1,108 @@ +"""Run upstream Pixal3D SLatFlowModel on CPU torch to produce a reference output. + +The SLAT flows contain no sparse convolution, so upstream is runnable here once the two +flash-attn kernels are swapped for torch SDPA. That gives a real numerical oracle. +""" +import json +import sys + +import numpy as np +import torch +import torch.nn.functional as F + +sys.path.insert(0, "/Users/m3ultra/Documents/pixal3d_mrp_mlx/upstream/Pixal3D") + +import pixal3d.modules.sparse as sp # noqa: E402 +from pixal3d.modules.sparse.basic import VarLenTensor # noqa: E402 +import pixal3d.modules.attention.modules as AM # noqa: E402 +import pixal3d.modules.sparse.attention.modules as SAM # noqa: E402 + + +def _dense_sdpa(q, k, v, **kw): + return F.scaled_dot_product_attention( + q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3) + ).permute(0, 2, 1, 3) + + +AM.scaled_dot_product_attention = _dense_sdpa + + +def _flatten_kv(t): + """-> (feats [T,H,C], list of (start,stop)). Handles varlen and dense [N,L,H,C].""" + if isinstance(t, VarLenTensor): + return t.feats, [(s.start, s.stop) for s in t.layout] + if t.dim() == 4: # [N, L, H, C] dense -> flatten, one segment per batch item + n, l = t.shape[0], t.shape[1] + return t.reshape(n * l, *t.shape[2:]), [(i * l, (i + 1) * l) for i in range(n)] + return t, [(0, t.shape[0])] + + +def _sparse_sdpa(*args): + if len(args) == 1: + ref = args[0] + f = ref.feats + q, k, v = f[:, 0], f[:, 1], f[:, 2] + lq = lkv = [(s.start, s.stop) for s in ref.layout] + elif len(args) == 2: + ref, kv = args + q, lq = _flatten_kv(ref) + kvf, lkv = _flatten_kv(kv) + k, v = kvf[:, 0], kvf[:, 1] + else: + ref, k_, v_ = args + q, lq = _flatten_kv(ref) + k, lkv = _flatten_kv(k_) + v, _ = _flatten_kv(v_) + outs = [] + for (a, b), (c, d) in zip(lq, lkv): + o = F.scaled_dot_product_attention( + q[a:b].permute(1, 0, 2)[None], + k[c:d].permute(1, 0, 2)[None], + v[c:d].permute(1, 0, 2)[None], + ) + outs.append(o[0].permute(1, 0, 2)) + out = torch.cat(outs, 0) + return ref.replace(out) if isinstance(ref, VarLenTensor) else out + + +SAM.sparse_scaled_dot_product_attention = _sparse_sdpa + +from pixal3d.models.structured_latent_flow import SLatFlowModel # noqa: E402 +from safetensors.torch import load_file # noqa: E402 + +CK = ( + "/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts/" + "slat_flow_img2shape_dit_1_3B_512_bf16" +) +cfg = json.load(open(CK + ".json")) +a = dict(cfg["args"]) +a["dtype"] = "float32" +m = SLatFlowModel(**a).eval() +m.load_state_dict( + {k: v.float() for k, v in load_file(CK + ".safetensors").items()}, strict=False +) + +rng = np.random.default_rng(0) +N = 64 +co = np.concatenate( + [np.zeros((N, 1), dtype=np.int32), rng.integers(0, 32, (N, 3)).astype(np.int32)], 1 +) +co = co[np.lexsort((co[:, 3], co[:, 2], co[:, 1], co[:, 0]))] +fe = rng.standard_normal((N, 32)).astype(np.float32) +g = (rng.standard_normal((1, 37, 1024)) * 0.1).astype(np.float32) +pr = (rng.standard_normal((N, 2048)) * 0.1).astype(np.float32) +t = torch.tensor(np.array([0.5], dtype=np.float32)) + +with torch.no_grad(): + out = m( + sp.SparseTensor(torch.tensor(fe), torch.tensor(co)), + t, + { + "global": torch.tensor(g), + "proj": sp.SparseTensor(torch.tensor(pr), torch.tensor(co)), + }, + ) +o = out.feats.numpy() +print(f" UPSTREAM slat: {o.shape} mean {o.mean():+.5f} std {o.std():.5f}") +for n, v in [("ref", o), ("co", co), ("fe", fe), ("g", g), ("pr", pr)]: + np.save(f"/tmp/px_slat_{n}.npy", v)