trellis-2-mrp-mlx/mlx_backend/vae_decoders.py

"""
VAE decoders in MLX: SparseUnetVaeDecoder and FlexiDualGridVaeDecoder.
"""
import logging
import time
import mlx.core as mx
import mlx.nn as nn
from typing import List, Optional, Tuple
from .norm import LayerNorm32
from .sparse_tensor import MlxSparseTensor
from .sparse_conv import MlxSparseConv3d
from .vae_blocks import (
    MlxSparseLinear,
    MlxSparseConvNeXtBlock3d,
    MlxSparseResBlockC2S3d,
    MlxSparseResBlockUpsample3d,
)

logger = logging.getLogger(__name__)


class MlxSparseUnetVaeDecoder(nn.Module):
    """
    Sparse UNet VAE decoder.
    Matches PyTorch SparseUnetVaeDecoder architecture.
    """

    def __init__(
        self,
        out_channels: int,
        model_channels: list,
        latent_channels: int,
        num_blocks: list,
        block_type: list,
        up_block_type: list,
        block_args: list = None,
        use_fp16: bool = False,
        pred_subdiv: bool = True,
    ):
        super().__init__()
        self.out_channels = out_channels
        self.model_channels = model_channels
        self.pred_subdiv = pred_subdiv
        self.use_fp16 = use_fp16

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

        # Build blocks
        BLOCK_TYPES = {
            'SparseConvNeXtBlock3d': MlxSparseConvNeXtBlock3d,
        }
        UP_BLOCK_TYPES = {
            'SparseResBlockC2S3d': MlxSparseResBlockC2S3d,
            'SparseResBlockUpsample3d': MlxSparseResBlockUpsample3d,
        }

        self.blocks = []
        for i in range(len(num_blocks)):
            stage = []
            BlockClass = BLOCK_TYPES[block_type[i]]
            for j in range(num_blocks[i]):
                stage.append(BlockClass(model_channels[i]))
            if i < len(num_blocks) - 1:
                UpBlockClass = UP_BLOCK_TYPES[up_block_type[i]]
                stage.append(UpBlockClass(
                    model_channels[i], model_channels[i + 1],
                    pred_subdiv=pred_subdiv,
                ))
            self.blocks.append(stage)

    def __call__(
        self,
        x: MlxSparseTensor,
        guide_subs: list = None,
        return_subs: bool = False,
    ):
        logger.info("[MLX] VAE decoder forward: N=%d, latent_ch=%d", x.feats.shape[0], x.feats.shape[1])
        t_total = time.time()

        h = self.from_latent(x)
        if self.use_fp16:
            h = h.replace(h.feats.astype(mx.float16))

        subs = []
        for i, stage in enumerate(self.blocks):
            t_stage = time.time()
            for j, block in enumerate(stage):
                if i < len(self.blocks) - 1 and j == len(stage) - 1:
                    # Up block
                    if self.pred_subdiv:
                        h, sub = block(h)
                        subs.append(sub)
                    else:
                        h = block(h, subdiv=guide_subs[i] if guide_subs is not None else None)
                else:
                    h = block(h)

            # Eval once per stage (not per block) to bound memory
            mx.eval(h.feats)
            dt_stage = time.time() - t_stage
            logger.info("[MLX] VAE stage %d/%d: N=%d, channels=%d, %.2fs",
                        i + 1, len(self.blocks), h.feats.shape[0], h.feats.shape[1], dt_stage)

        h = h.replace(h.feats.astype(x.feats.dtype))

        # Final layer norm (two-pass for parity)
        feats = h.feats.astype(mx.float32)
        mean = mx.mean(feats, axis=-1, keepdims=True)
        centered = feats - mean
        var = mx.mean(centered * centered, axis=-1, keepdims=True)
        feats = centered * mx.rsqrt(var + 1e-5)
        h = h.replace(feats.astype(x.feats.dtype))

        h = self.output_layer(h)

        dt_total = time.time() - t_total
        logger.info("[MLX] VAE decoder done: N=%d, out_ch=%d, total %.2fs",
                    h.feats.shape[0], h.feats.shape[1], dt_total)

        if return_subs:
            return h, subs
        return h

    def upsample(self, x: MlxSparseTensor, upsample_times: int) -> mx.array:
        """Run decoder up to upsample_times stages, return upsampled coords."""
        logger.info("[MLX] VAE upsample: N=%d, upsample_times=%d", x.feats.shape[0], upsample_times)
        h = self.from_latent(x)
        if self.use_fp16:
            h = h.replace(h.feats.astype(mx.float16))

        for i, stage in enumerate(self.blocks):
            if i == upsample_times:
                return h.coords
            for j, block in enumerate(stage):
                if i < len(self.blocks) - 1 and j == len(stage) - 1:
                    h, sub = block(h)
                else:
                    h = block(h)
                mx.eval(h.feats)

        return h.coords


class MlxFlexiDualGridVaeDecoder(nn.Module):
    """
    FlexiDualGrid VAE decoder — wraps SparseUnetVaeDecoder.
    Outputs mesh vertices + intersection logits + quad_lerp.
    """

    def __init__(
        self,
        resolution: int,
        model_channels: list,
        latent_channels: int,
        num_blocks: list,
        block_type: list,
        up_block_type: list,
        block_args: list = None,
        use_fp16: bool = False,
        voxel_margin: float = 0.5,
    ):
        super().__init__()
        self.resolution = resolution
        self.voxel_margin = voxel_margin

        # out_channels = 7 (3 vertex + 3 intersection + 1 quad_lerp)
        self.decoder = MlxSparseUnetVaeDecoder(
            out_channels=7,
            model_channels=model_channels,
            latent_channels=latent_channels,
            num_blocks=num_blocks,
            block_type=block_type,
            up_block_type=up_block_type,
            block_args=block_args,
            use_fp16=use_fp16,
            pred_subdiv=True,
        )

    def set_resolution(self, resolution: int):
        self.resolution = resolution

    def __call__(self, x: MlxSparseTensor, return_subs: bool = False, **kwargs):
        decoded = self.decoder(x, return_subs=return_subs, **kwargs)

        if return_subs:
            h, subs = decoded
        else:
            h = decoded
            subs = None

        # Post-process: vertices = sigmoid(h[:, :3]), intersected = h[:, 3:6] > 0
        feats = h.feats
        vertices_feats = (1 + 2 * self.voxel_margin) * mx.sigmoid(feats[:, :3]) - self.voxel_margin
        intersected_feats = (feats[:, 3:6] > 0).astype(feats.dtype)
        quad_lerp_feats = mx.where(feats[:, 6:7] > 0, feats[:, 6:7], mx.zeros_like(feats[:, 6:7])) + mx.log1p(mx.exp(-mx.abs(feats[:, 6:7])))  # softplus

        result = (h, vertices_feats, intersected_feats, quad_lerp_feats)

        if return_subs:
            return result, subs
        return result

    def upsample(self, x: MlxSparseTensor, upsample_times: int) -> mx.array:
        return self.decoder.upsample(x, upsample_times)