8faf6d592a
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.
193 lines
6.3 KiB
Python
193 lines
6.3 KiB
Python
"""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
|