pixal3d_mrp_mlx/pixal3d_mlx/sampler.py

"""Flow-matching Euler sampler in MLX, with classifier-free guidance.

Rectified-flow sampling: the model predicts a velocity `v` and the update is simply

    x_{t-1} = x_t - (t - t_prev) * v

Timesteps run from 1 down to 0. `rescale_t` warps the schedule
(`r*t / (1 + (r-1)*t)`), which concentrates steps near t=1 for r>1 — Pixal3D uses this
to spend more compute early where the structure is decided.

Two guidance behaviours, both from upstream:

* **CFG** — `pred = g*pos + (1-g)*neg`. Note this is a LERP, not the more common
  `neg + g*(pos - neg)`; they agree only because upstream's `g` is defined on the same
  scale. Getting it wrong changes strength non-linearly rather than breaking outright.
* **Guidance interval** — guidance applies only while `lo <= t <= hi`; outside that
  window strength is forced to 1 (i.e. conditional only, one model call instead of two).
  That is a real speedup as well as a quality choice.

Works on both dense `mx.array` latents and `SparseTensor` ones, since the SLAT stage
denoises sparse features — hence the small `_lift`/`_like` helpers rather than raw
arithmetic.
"""

from __future__ import annotations

from typing import Any, Callable, List, Optional, Tuple

import mlx.core as mx

from trellis_sparse_mlx import SparseTensor


def _feats(x):
    return x.feats if isinstance(x, SparseTensor) else x


def _like(ref, feats):
    """Rebuild `ref`'s container around new features."""
    return ref.replace(feats) if isinstance(ref, SparseTensor) else feats


class FlowEulerSampler:
    def __init__(self, sigma_min: float = 1e-5):
        self.sigma_min = sigma_min

    # -- conversions between the model's velocity and x0/eps -------------------
    def v_to_xstart_eps(self, x_t, t: float, v):
        xf, vf = _feats(x_t), _feats(v)
        eps = (1 - t) * vf + xf
        x0 = (1 - self.sigma_min) * xf - (self.sigma_min + (1 - self.sigma_min) * t) * vf
        return x0, eps

    def pred_to_xstart(self, x_t, t: float, pred):
        xf, pf = _feats(x_t), _feats(pred)
        return (1 - self.sigma_min) * xf - (
            self.sigma_min + (1 - self.sigma_min) * t
        ) * pf

    def xstart_to_pred(self, x_t, t: float, x_0):
        """Inverse of `pred_to_xstart` — needed by the CFG rescale."""
        xf = _feats(x_t)
        return ((1 - self.sigma_min) * xf - _feats(x_0)) / (
            self.sigma_min + (1 - self.sigma_min) * t
        )

    # -- model call ------------------------------------------------------------
    def _call(self, model: Callable, x_t, t: float, cond) -> Any:
        b = _feats(x_t).shape[0] if not isinstance(x_t, SparseTensor) else len(x_t.layout)
        tt = mx.full((b,), 1000.0 * t, dtype=mx.float32)  # upstream scales t by 1000
        return model(x_t, tt, cond)

    def _inference(
        self,
        model,
        x_t,
        t: float,
        cond,
        neg_cond=None,
        guidance_strength: float = 1.0,
        guidance_interval: Optional[Tuple[float, float]] = None,
        guidance_rescale: float = 0.0,
    ):
        g = guidance_strength
        if guidance_interval is not None and not (
            guidance_interval[0] <= t <= guidance_interval[1]
        ):
            g = 1.0  # outside the window: conditional only, and only one model call
        if g == 1.0 or neg_cond is None:
            return _feats(self._call(model, x_t, t, cond))
        if g == 0.0:
            return _feats(self._call(model, x_t, t, neg_cond))
        pos = _feats(self._call(model, x_t, t, cond))
        neg = _feats(self._call(model, x_t, t, neg_cond))
        pred = g * pos + (1 - g) * neg  # LERP, matching upstream

        # CFG rescale (Lin et al., "Common Diffusion Noise Schedules ... are Flawed").
        # High guidance inflates the variance of x0; this pulls it back to the
        # conditional branch's std. NOT optional here — the shipped ss config sets
        # guidance_rescale=0.7 and shape_slat 0.5, so omitting it silently overcooks
        # every structure prediction at the pipeline's own default settings.
        if guidance_rescale > 0:
            x0_pos = self.pred_to_xstart(x_t, t, pos)
            x0_cfg = self.pred_to_xstart(x_t, t, pred)
            axes = tuple(range(1, x0_pos.ndim))
            std_pos = mx.sqrt(mx.var(x0_pos, axis=axes, keepdims=True, ddof=1))
            std_cfg = mx.sqrt(mx.var(x0_cfg, axis=axes, keepdims=True, ddof=1))
            x0 = guidance_rescale * (x0_cfg * (std_pos / std_cfg)) + (1 - guidance_rescale) * x0_cfg
            pred = self.xstart_to_pred(x_t, t, x0)
        return pred

    # -- loop ------------------------------------------------------------------
    @staticmethod
    def timesteps(steps: int, rescale_t: float = 1.0) -> List[Tuple[float, float]]:
        seq = [1.0 - i / steps for i in range(steps + 1)]
        seq = [rescale_t * t / (1 + (rescale_t - 1) * t) for t in seq]
        return [(seq[i], seq[i + 1]) for i in range(steps)]

    def sample(
        self,
        model,
        noise,
        cond,
        neg_cond=None,
        steps: int = 25,
        rescale_t: float = 1.0,
        guidance_strength: float = 1.0,
        guidance_interval: Optional[Tuple[float, float]] = None,
        guidance_rescale: float = 0.0,
        progress: Optional[Callable[[int, int], None]] = None,
    ):
        x = noise
        pairs = self.timesteps(steps, rescale_t)
        for i, (t, t_prev) in enumerate(pairs):
            v = self._inference(
                model, x, t, cond, neg_cond, guidance_strength, guidance_interval,
                guidance_rescale,
            )
            x = _like(x, _feats(x) - (t - t_prev) * v)
            mx.eval(_feats(x))
            if progress:
                progress(i + 1, steps)
        return x