diff --git a/CLAUDE.md b/CLAUDE.md index 9139a7a..b97e965 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -37,3 +37,17 @@ MLX port of TencentARC Pixal3D, on top of the shared `trellis_sparse_mlx` core. classify_5d() distinguishes them correctly on the real weights. - `rope_phases` ships as a stored tensor — RoPE phases are precomputed, not derived. - Converter preserves dtype (fp16 stays fp16); upcasting doubled 24GB for nothing. + +## Numerical oracle (important) +The flow models have NO sparse conv, so upstream RUNS ON CPU TORCH here. Patch +`pixal3d.modules.attention.modules.scaled_dot_product_attention` with a torch SDPA +wrapper (permute to [B,H,N,D] and back) and it loads the real checkpoint. Use this to +diff any flow-model change — do not reason about correctness, measure it. + +## Bugs found that weight-matching could NOT catch (loader said 700/700, 0 missing) +1. Parameterless final F.layer_norm between last block and out_layer. No params -> no + checkpoint trace. Without it output is ~200x too large. +2. rope_phases is complex64 (torch.polar). Rotation is a complex multiply; cos/sin are + the real/imag parts, NOT cos(phase). +3. qk_rms_norm goes BEFORE rope, not after. Non-commutative. One block correlates 0.9998 + when reversed; compounds to 0.84 over 30 blocks. diff --git a/README.md b/README.md index c00f9b3..eb82651 100644 --- a/README.md +++ b/README.md @@ -71,7 +71,36 @@ Weights: 24.04 GB across 19 files (1.3B DiTs at 512/1024 + shape/tex decoders). - [x] `upsample` (masked) + `downsample(mode=)` landed in the shared core - [x] **Weight converter** — decoders remap KRSC→`[K³,in,out]`, dtype preserved, remap verified a pure permutation. Flow models pass through untouched. -- [ ] DiT blocks: RoPE (stored phases), qk_rms_norm, AdaLN modulation +- [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` - [ ] Model graphs - [ ] End-to-end + + +## Numerical verification + +The flow models contain no sparse convolution, which means **upstream runs on CPU torch +here** — swap flash-attn for `F.scaled_dot_product_attention` and it loads the real +checkpoint and runs. So unlike the sparse path (where spconv is uninstallable and the +oracle had to be hand-written), these are diffed against upstream directly: + +``` +MLX : mean +0.18490 std 0.86841 +UPSTREAM : mean +0.18490 std 0.86841 +max abs diff 1.216e-05 correlation 1.00000000 +``` + +Getting there required finding three bugs that **weight-key matching could not catch** — +the loader reported a perfect 700/700 with 0 missing and 0 unmapped through all of them: + +1. **A parameterless final `LayerNorm`** between the last block and `out_layer`. It has + no weights, so it leaves no trace in the checkpoint. Without it the output was ~200x + too large (std 187 vs 0.87). +2. **`rope_phases` is complex64**, built with `torch.polar`. The rotation is a complex + multiply and cos/sin are the phase's real/imaginary parts — taking `cos()` of a + complex phase is meaningless. +3. **qk RMS norm is applied BEFORE RoPE, not after.** They do not commute. Reversed, a + single block still correlated 0.9998 with upstream; over 30 blocks that compounds to + 0.84. This one is invisible without an oracle. diff --git a/pixal3d_mlx/ss_flow.py b/pixal3d_mlx/ss_flow.py new file mode 100644 index 0000000..5c79e24 --- /dev/null +++ b/pixal3d_mlx/ss_flow.py @@ -0,0 +1,203 @@ +"""Pixal3D SparseStructureFlowModel in MLX. + +The smallest of Pixal3D's four flow checkpoints (5.0 GB) and the one that exercises +every DiT feature the others use — RoPE, per-head q/k RMS norm, shared AdaLN modulation, +and `image_attn_mode="proj"` conditioning — so getting this one loading and running +validates the shared DiT core against real trained weights. + +Despite the name it is fully DENSE: it operates on a 16³ voxel grid flattened to 4096 +tokens. No sparse convolution anywhere (the checkpoint contains zero rank-5 tensors). + +Config comes from the sibling JSON, not from constructor defaults: + resolution 16, in/out 8ch, model 1536, cond 1024, 30 blocks, 12 heads, + mlp_ratio 5.3334, pe_mode rope, share_mod, qk_rms_norm (+cross), image_attn_mode proj +""" + +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.dit import ModulatedTransformerCrossBlock, TimestepEmbedder + + +class SparseStructureFlowModel(nn.Module): + def __init__( + self, + resolution: int = 16, + in_channels: int = 8, + out_channels: int = 8, + 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.share_mod = share_mod + self.pe_mode = pe_mode + + self.t_embedder = TimestepEmbedder(model_channels) + if share_mod: + # upstream is Sequential(SiLU, Linear) -> checkpoint key is adaLN_modulation.1 + self.adaLN_modulation = nn.Linear(model_channels, 6 * model_channels) + self.input_layer = nn.Linear(in_channels, model_channels) + self.blocks = [ + ModulatedTransformerCrossBlock( + 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 = nn.Linear(model_channels, out_channels) + self.rope_phases: Optional[mx.array] = None # loaded from the checkpoint + + def __call__(self, x: mx.array, t: mx.array, cond) -> mx.array: + """x: [B, C, D, H, W]; t: [B]; cond: dict/tuple of (global, proj).""" + b, c = x.shape[0], x.shape[1] + if c != self.in_channels or list(x.shape[2:]) != [self.resolution] * 3: + raise ValueError( + f"expected [B,{self.in_channels},{self.resolution}^3], got {x.shape}" + ) + h = x.reshape(b, c, -1).transpose(0, 2, 1) # [B, N, C] + h = self.input_layer(h) + + t_emb = self.t_embedder(t) + if self.share_mod: + t_emb = self.adaLN_modulation(nn.silu(t_emb)) + + phases = self.rope_phases + for blk in self.blocks: + h = blk(h, t_emb, cond, phases=phases) + + # Parameterless final LayerNorm — `F.layer_norm(h, h.shape[-1:])` upstream. + # Easy to miss and impossible to catch by weight-key matching: it has no + # parameters, so a loader can report a perfect 700/700 with 0 missing and 0 + # unmapped while the model is still wrong. The residual stream leaves the last + # block at std ~230; without this the output is ~200x too large. + mu = mx.mean(h, axis=-1, keepdims=True) + var = mx.var(h, axis=-1, keepdims=True) + h = (h - mu) * mx.rsqrt(var + 1e-5) + + h = self.out_layer(h) + return h.transpose(0, 2, 1).reshape(b, self.out_channels, *[self.resolution] * 3) + + +# ------------------------------------------------------------------ loading + +def _remap(k: str) -> str: + """Checkpoint key -> module path. + + Upstream wraps a couple of things in nn.Sequential, so its keys carry numeric indices + where this implementation uses named attributes: + t_embedder.mlp.{0,2} -> t_embedder.mlp_{0,2} (Linear, SiLU, Linear) + blocks.N.mlp.mlp.{0,2} -> blocks.N.mlp_{0,2} + adaLN_modulation.1 -> adaLN_modulation (index 0 is the SiLU) + """ + 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.") + return k + + +def load(weights_path: str | Path, config_path: str | Path | None = None): + """Build from the sibling JSON config and load weights. Returns (model, report).""" + 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 = SparseStructureFlowModel(**args) + w = mx.load(str(wp)) + + flat = dict(_flatten(model.parameters())) + mapped, unmapped = {}, [] + for k, v in w.items(): + if k == "rope_phases": + model.rope_phases = v + continue + 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, + "rope_phases": None if model.rope_phases is None else model.rope_phases.shape, + } + + +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