mlx_backend: pure-MLX FlowEuler+CFG(+interval) sampler loop

Replaces the torch CPU sampler + per-forward adapter bounces: one
conversion per stage boundary, batched dense CFG, once-per-stage
concat_cond, one eval/step. Env-gated (TRELLIS2_MLX_SAMPLER=0 reverts;
TRELLIS2_MLX_COMPILE=1 compiles the step forward). Parity: raw mesh
within 0.03% of torch-sampler baseline at seed 42. Sampling 56.8->54.4s
(m3ultra) — honest ~4%; the win is architectural (stable shapes for
compile, no host round-trips) not headline speed.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

mlx_backend/mlx_samplers.py

@@ -0,0 +1,159 @@
+"""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)})

mlx_backend/pipeline.py

@@ -249,6 +249,18 @@ def create_mlx_pipeline(
     )(**args['tex_slat_sampler']['args'])
     pipeline.tex_slat_sampler_params = args['tex_slat_sampler']['params']
 
+    # Pure-MLX sampler loop (rung-3): kills per-forward host bounces + syncs.
+    # Rollback: TRELLIS2_MLX_SAMPLER=0 keeps the upstream torch samplers.
+    import os as _os
+    if _os.environ.get('TRELLIS2_MLX_SAMPLER', '1') == '1':
+        from .mlx_samplers import MlxFlowEulerGuidanceIntervalSampler
+        for _stage in ('sparse_structure_sampler', 'shape_slat_sampler',
+                       'tex_slat_sampler'):
+            if args[_stage]['name'] == 'FlowEulerGuidanceIntervalSampler':
+                setattr(pipeline, _stage,
+                        MlxFlowEulerGuidanceIntervalSampler(**args[_stage]['args']))
+        print('[MLX] Pure-MLX sampler loop active (TRELLIS2_MLX_SAMPLER=0 to revert)')
+
     # Normalization
     pipeline.shape_slat_normalization = args['shape_slat_normalization']
     pipeline.tex_slat_normalization = args['tex_slat_normalization']