"""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, concat_cond: SparseTensor | None = None) -> SparseTensor: # The texture flow is `imgshape2tex`: it denoises 32 PBR channels while SEEING # the shape latent, so in_channels is 64 against out_channels 32. Upstream does # `sparse_cat([x, concat_cond], dim=-1)`; both share coords, so it reduces to a # plain channel concat. Without it the input layer gets half its expected width. if concat_cond is not None: x = x.replace(mx.concatenate([x.feats, concat_cond.feats], axis=-1)) 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