40688778c2
700/700 params, max abs diff 1.2e-5 on the real 1.3B checkpoint. Key realisation: the flow models have no sparse conv, so upstream RUNS on CPU torch with flash-attn swapped for SDPA. That gives a real numerical oracle - unavailable for the sparse path, where spconv cannot be installed at all. It was needed. Three bugs survived a loader reporting a perfect 700/700 with zero missing and zero unmapped keys: - a parameterless final LayerNorm (no params -> no checkpoint trace) that the output was 200x too large without - rope_phases being complex64, so the rotation is a complex multiply - qk_rms_norm belonging before rope rather than after Weight-key matching is necessary but nowhere near sufficient for a port.
204 lines
7.5 KiB
Python
204 lines
7.5 KiB
Python
"""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
|