pixal3d_mrp_mlx/pixal3d_mlx/slat_flow.py

"""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