pixal3d_mrp_mlx/pixal3d_mlx/decoders.py

"""Pixal3D sparse VAE decoders in MLX — `shape_dec` and `tex_dec`.

Both are the same U-Net: stages of `SparseConvNeXtBlock3d`, each stage ending in a
`SparseResBlockC2S3d` that upsamples via channel->spatial. The occupied set grows
selectively at every stage, driven by a predicted subdivision mask, which is how the
decoder reaches 256^3-equivalent detail without ever materialising a dense grid.

    model_channels [1024, 512, 256, 128, 64]
    num_blocks     [4, 16, 8, 4, 0]        <- last stage is empty and has no up-block
    block_type     SparseConvNeXtBlock3d
    up_block_type  SparseResBlockC2S3d

`FlexiDualGridVaeDecoder` (shape_dec) subclasses it and hardcodes **7** output channels
regardless of config — 3 for vertex offsets (squashed through a sigmoid widened by
`voxel_margin`) plus 4 more — so its `out_channels` is not a config field at all.
`SparseUnetVaeDecoder` (tex_dec) uses the config's `out_channels` (6) and sets
`pred_subdiv=False`, meaning it needs subdivision masks supplied from outside rather
than predicting its own.

NOTE ON VERIFICATION: unlike the flow models and ss_dec, these use sparse convolution,
so upstream cannot run here (no spconv on Metal) and there is no numerical oracle. They
are built from blocks that ARE individually tested, but the assembled graphs are checked
only for shape, finiteness and plausible behaviour. Treat them as less certain than
anything reporting a correlation.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Dict, List, Optional

import mlx.core as mx
import mlx.nn as nn

from trellis_sparse_mlx import (
    SparseConvNeXtBlock3d,
    SparseLinear,
    SparseResBlockC2S3d,
    SparseTensor,
)

_BLOCKS = {
    "SparseConvNeXtBlock3d": SparseConvNeXtBlock3d,
    "SparseResBlockC2S3d": SparseResBlockC2S3d,
}


class SparseUnetVaeDecoder(nn.Module):
    def __init__(
        self,
        out_channels: int,
        model_channels: List[int],
        latent_channels: int,
        num_blocks: List[int],
        block_type: List[str],
        up_block_type: List[str],
        block_args: Optional[List[Dict[str, Any]]] = None,
        pred_subdiv: bool = True,
        **_ignored,
    ):
        super().__init__()
        self.out_channels = out_channels
        self.model_channels = list(model_channels)
        self.num_blocks = list(num_blocks)
        self.pred_subdiv = pred_subdiv
        block_args = block_args or [{} for _ in num_blocks]

        self.from_latent = SparseLinear(latent_channels, model_channels[0])
        self.output_layer = SparseLinear(model_channels[-1], out_channels)

        stages = []
        for i in range(len(num_blocks)):
            stage = [
                _BLOCKS[block_type[i]](model_channels[i], **block_args[i])
                for _ in range(num_blocks[i])
            ]
            if i < len(num_blocks) - 1:
                stage.append(
                    _BLOCKS[up_block_type[i]](
                        model_channels[i],
                        model_channels[i + 1],
                        pred_subdiv=pred_subdiv,
                    )
                )
            stages.append(stage)
        self.blocks = stages

    def __call__(
        self,
        x: SparseTensor,
        guide_subs: Optional[List[SparseTensor]] = None,
        return_subs: bool = False,
    ):
        h = self.from_latent(x)
        subs = []
        for i, stage in enumerate(self.blocks):
            for j, blk in enumerate(stage):
                is_up = i < len(self.blocks) - 1 and j == len(stage) - 1
                if is_up:
                    if self.pred_subdiv:
                        h, sub = blk(h)
                        subs.append(sub)
                    else:
                        h = blk(h, subdiv=guide_subs[i] if guide_subs else None)
                else:
                    h = blk(h)

        # Same parameterless final LayerNorm as every other model in this family.
        f = h.feats
        f = (f - mx.mean(f, -1, keepdims=True)) * mx.rsqrt(
            mx.var(f, -1, keepdims=True) + 1e-5
        )
        out = self.output_layer(h.replace(f))
        return (out, subs) if return_subs else out

    def upsample(self, x: SparseTensor, upsample_times: int) -> mx.array:
        """Run only the first `upsample_times` stages and return the COORDS.

        The cascade needs a refined coordinate set, not features: the low-res SLAT is
        pushed part-way through the decoder purely so its predicted subdivisions grow
        the occupied set, and the resulting coords seed the high-res flow. Each stage
        doubles resolution, so `upsample_times=4` yields coords at 16x the input.

        Deliberately stops BEFORE running stage `upsample_times`, matching upstream —
        running one stage further would return coords at 2x the intended resolution
        and silently misplace every voxel in the high-res pass.
        """
        if not self.pred_subdiv:
            raise ValueError("upsample needs a decoder with pred_subdiv=True")
        h = self.from_latent(x)
        for i, stage in enumerate(self.blocks):
            if i == upsample_times:
                return h.coords
            for j, blk in enumerate(stage):
                if i < len(self.blocks) - 1 and j == len(stage) - 1:
                    h, _ = blk(h)
                else:
                    h = blk(h)
        return h.coords


class FlexiDualGridVaeDecoder(SparseUnetVaeDecoder):
    """shape_dec. Emits 7 channels; the first 3 are vertex offsets within the voxel."""

    def __init__(
        self,
        resolution: int,
        model_channels: List[int],
        latent_channels: int,
        num_blocks: List[int],
        block_type: List[str],
        up_block_type: List[str],
        block_args: Optional[List[Dict[str, Any]]] = None,
        voxel_margin: float = 0.5,
        **_ignored,
    ):
        # out_channels is hardcoded to 7 upstream, not read from config
        super().__init__(
            7,
            model_channels,
            latent_channels,
            num_blocks,
            block_type,
            up_block_type,
            block_args,
            pred_subdiv=True,
        )
        self.resolution = resolution
        self.voxel_margin = voxel_margin

    def decode_vertices(self, h: SparseTensor) -> mx.array:
        """Channels 0:3 -> vertex offsets, sigmoid widened by `voxel_margin`.

        `(1 + 2m)*sigmoid(v) - m` maps to [-m, 1+m], letting a vertex sit slightly
        outside its own voxel — that overshoot is what lets the dual grid represent
        surfaces that do not pass exactly through cell centres.
        """
        m = self.voxel_margin
        return (1 + 2 * m) * mx.sigmoid(h.feats[..., 0:3]) - m


# ------------------------------------------------------------------ loading

def _remap(k: str) -> str:
    """Upstream: blocks.<stage>.<idx>.<...>; SparseLinear wraps an nn.Linear."""
    for name in ("from_latent", "output_layer", "to_subdiv"):
        k = k.replace(f"{name}.weight", f"{name}.linear.weight")
        k = k.replace(f"{name}.bias", f"{name}.linear.bias")
    k = k.replace(".mlp.0.", ".mlp_0.").replace(".mlp.2.", ".mlp_2.")
    k = k.replace(".conv1.conv.", ".conv1.").replace(".conv2.conv.", ".conv2.")
    k = k.replace(".conv.conv.", ".conv.")
    k = k.replace(".norm.weight", ".norm.weight").replace(".norm1.weight", ".norm1.weight")
    return k


def load(weights_path: str | Path, config_path: str | Path):
    cfg = json.loads(Path(config_path).read_text())
    name, args = cfg["name"], dict(cfg.get("args", {}))
    args.pop("use_fp16", None)
    cls = {
        "FlexiDualGridVaeDecoder": FlexiDualGridVaeDecoder,
        "SparseUnetVaeDecoder": SparseUnetVaeDecoder,
    }[name]
    model = cls(**args)

    w = mx.load(str(weights_path))
    flat = dict(_flatten(model.parameters()))
    mapped, unmapped, mismatched = {}, [], []
    for k, v in w.items():
        m = _remap(k)
        if m in flat:
            if flat[m].shape != v.shape:
                mismatched.append(f"{k} -> {m}: want {flat[m].shape} got {v.shape}")
                continue
            mapped[m] = v
        else:
            unmapped.append(f"{k} -> {m} {tuple(v.shape)}")
    missing = [k for k in flat if k not in mapped]
    if mapped:
        model.update(_unflatten(mapped))
    return model, {
        "config": name,
        "loaded": len(mapped),
        "params": len(flat),
        "missing": missing,
        "unmapped": unmapped,
        "mismatched": mismatched,
    }


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