"""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)})