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

"""Pure-MLX FlowEuler + CFG(+interval) sampler for the mlx-experimental backend.

Replaces the torch CPU sampler loop and its per-call adapter host bounces:
- tensors convert ONCE at stage entry/exit (was: every forward, 48+/stage)
- dense CFG pairs run as ONE batched forward (was: two synced forwards)
- `concat_cond` converts once per stage (was: full SparseTensor rebuild/step)
- one mx.eval per STEP (was: per forward), keeping the graph on-device
- TRELLIS2_MLX_COMPILE=1 additionally mx.compile's the step forward
  (fixed shapes per stage; corridorkey/Hunyuan-T1 pattern)

Numerics mirror trellis2/pipelines/samplers/{flow_euler.py,
classifier_free_guidance_mixin.py, guidance_interval_mixin.py} exactly
(inclusive interval on t; torch-default unbiased std in CFG rescale).

Rollback: TRELLIS2_MLX_SAMPLER=0 keeps the upstream torch samplers.
"""
import os
from typing import Any, Optional

import mlx.core as mx
import numpy as np
import torch
from easydict import EasyDict as edict
from tqdm import tqdm

from .adapters import _mlx_to_sparse, _mx_to_torch, _sparse_to_mlx, _torch_to_mx
from .sparse_tensor import MlxSparseTensor


class MlxFlowEulerGuidanceIntervalSampler:
    """MLX drop-in for FlowEulerGuidanceIntervalSampler (same .sample API)."""

    def __init__(self, sigma_min: float):
        self.sigma_min = sigma_min

    # ---- math mirrors (flow_euler.py) ----
    def _pred_to_xstart(self, x_t, t, pred):
        return (1 - self.sigma_min) * x_t - \
            (self.sigma_min + (1 - self.sigma_min) * t) * pred

    def _xstart_to_pred(self, x_t, t, x_0):
        return ((1 - self.sigma_min) * x_t - x_0) / \
            (self.sigma_min + (1 - self.sigma_min) * t)

    def _cfg_mix(self, x_t, t, v_pos, v_neg, g, g_rescale, feat_axes):
        pred = g * v_pos + (1 - g) * v_neg
        if g_rescale > 0:
            x_0_pos = self._pred_to_xstart(x_t, t, v_pos)
            x_0_cfg = self._pred_to_xstart(x_t, t, pred)
            std_pos = mx.sqrt(mx.var(x_0_pos, axis=feat_axes,
                                     keepdims=True, ddof=1))
            std_cfg = mx.sqrt(mx.var(x_0_cfg, axis=feat_axes,
                                     keepdims=True, ddof=1))
            x_0_res = x_0_cfg * (std_pos / std_cfg)
            x_0 = g_rescale * x_0_res + (1 - g_rescale) * x_0_cfg
            pred = self._xstart_to_pred(x_t, t, x_0)
        return pred

    @staticmethod
    def _t_pairs(steps, rescale_t):
        t_seq = np.linspace(1, 0, steps + 1)
        t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
        t_seq = t_seq.tolist()
        return [(t_seq[i], t_seq[i + 1]) for i in range(steps)]

    @staticmethod
    def _maybe_compile(fn):
        if os.environ.get("TRELLIS2_MLX_COMPILE", "0") == "1":
            return mx.compile(fn)
        return fn

    # ---- public API (matches upstream sampler.sample) ----
    def sample(self, model, noise, cond=None, neg_cond=None, steps: int = 50,
               rescale_t: float = 1.0, guidance_strength: float = 7.5,
               guidance_rescale: float = 0.0, guidance_interval=(0.0, 1.0),
               verbose: bool = True, tqdm_desc: str = "Sampling", **kwargs):
        raw = getattr(model, "_mlx", model)  # unwrap MlxFlowModelAdapter
        if isinstance(noise, torch.Tensor):
            return self._sample_dense(
                raw, noise, cond, neg_cond, steps, rescale_t,
                guidance_strength, guidance_rescale, guidance_interval,
                verbose, tqdm_desc)
        return self._sample_sparse(
            raw, noise, cond, neg_cond, steps, rescale_t,
            guidance_strength, guidance_rescale, guidance_interval,
            verbose, tqdm_desc, **kwargs)

    # ---- dense (sparse-structure stage): batched CFG ----
    def _sample_dense(self, raw, noise, cond, neg_cond, steps, rescale_t,
                      g, g_rescale, g_interval, verbose, desc):
        x = _torch_to_mx(noise)
        B = x.shape[0]
        c = _torch_to_mx(cond)
        n = _torch_to_mx(neg_cond) if neg_cond is not None else None
        c2 = mx.concatenate([c, n], axis=0) if n is not None else None
        feat_axes = tuple(range(1, x.ndim))

        fwd = self._maybe_compile(lambda xx, tt, cc: raw(xx, tt, cc))

        for t, t_prev in tqdm(self._t_pairs(steps, rescale_t), desc=desc,
                              disable=not verbose):
            in_iv = g_interval[0] <= t <= g_interval[1]
            eff_g = g if in_iv else 1.0
            t_m = mx.full((B,), 1000.0 * t, dtype=mx.float32)
            if eff_g == 1.0 or n is None:
                pred = fwd(x, t_m, c)
            elif eff_g == 0.0:
                pred = fwd(x, t_m, n)
            else:
                x2 = mx.concatenate([x, x], axis=0)
                t2 = mx.concatenate([t_m, t_m], axis=0)
                v2 = fwd(x2, t2, c2)
                pred = self._cfg_mix(x, t, v2[:B], v2[B:], eff_g,
                                     g_rescale if in_iv else 0.0, feat_axes)
            x = x - (t - t_prev) * pred
            mx.eval(x)
        return edict({"samples": _mx_to_torch(x)})

    # ---- sparse (slat stages): once-per-stage conversion, cached tensor ----
    def _sample_sparse(self, raw, noise, cond, neg_cond, steps, rescale_t,
                       g, g_rescale, g_interval, verbose, desc, **kwargs):
        xs = _sparse_to_mlx(noise)          # coords fixed for the whole stage
        template = xs                        # .replace() keeps spatial caches
        c = _torch_to_mx(cond)
        n = _torch_to_mx(neg_cond) if neg_cond is not None else None
        cc = None
        if kwargs.get("concat_cond") is not None:
            cc = _sparse_to_mlx(kwargs["concat_cond"])  # ONCE (was per-step)

        def step_fwd(feats, t_m, cond_a):
            xt = template.replace(feats)
            out = raw(xt, t_m, cond_a, **({"concat_cond": cc} if cc is not None else {}))
            return out.feats

        fwd = self._maybe_compile(step_fwd)

        feats = xs.feats
        # batch=1 sparse: std over all feat elements (== upstream non-batch dims)
        feat_axes = tuple(range(feats.ndim))

        for t, t_prev in tqdm(self._t_pairs(steps, rescale_t), desc=desc,
                              disable=not verbose):
            in_iv = g_interval[0] <= t <= g_interval[1]
            eff_g = g if in_iv else 1.0
            t_m = mx.array([1000.0 * t], dtype=mx.float32)
            if eff_g == 1.0 or n is None:
                pred = fwd(feats, t_m, c)
            elif eff_g == 0.0:
                pred = fwd(feats, t_m, n)
            else:
                v_pos = fwd(feats, t_m, c)
                v_neg = fwd(feats, t_m, n)
                pred = self._cfg_mix(feats, t, v_pos, v_neg, eff_g,
                                     g_rescale if in_iv else 0.0, feat_axes)
            feats = feats - (t - t_prev) * pred
            mx.eval(feats)

        return edict({"samples": _mlx_to_sparse(template.replace(feats),
                                                ref=noise)})