From fa972aa34547032b84e739c2537f3e6c006677d9 Mon Sep 17 00:00:00 2001 From: m3ultra Date: Mon, 20 Jul 2026 03:01:57 +1000 Subject: [PATCH] mlx_backend: pure-MLX FlowEuler+CFG(+interval) sampler loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- mlx_backend/mlx_samplers.py | 159 ++++++++++++++++++++++++++++++++++++ mlx_backend/pipeline.py | 12 +++ 2 files changed, 171 insertions(+) create mode 100644 mlx_backend/mlx_samplers.py diff --git a/mlx_backend/mlx_samplers.py b/mlx_backend/mlx_samplers.py new file mode 100644 index 0000000..9604a42 --- /dev/null +++ b/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)}) diff --git a/mlx_backend/pipeline.py b/mlx_backend/pipeline.py index 2ba44c4..cf72c1e 100644 --- a/mlx_backend/pipeline.py +++ b/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']