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.
This commit is contained in:
John
parent
40688778c2
commit
8faf6d592a
@ -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
|
||||
|
||||
|
||||
192
pixal3d_mlx/slat_flow.py
Normal file
192
pixal3d_mlx/slat_flow.py
Normal file
@ -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
|
||||
108
tests/oracle_slat.py
Normal file
108
tests/oracle_slat.py
Normal file
@ -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)
|
||||
Loading…
Reference in New Issue
Block a user