d43a6dfd69
The structure stage now runs end to end on Metal: noise -> ss_flow (12-step Euler) -> latent -> ss_dec -> 64^3 occupancy grid. 759 ms/step for the 1.3B DiT, 50 ms for the decoder. Output lands at 1.36% occupancy, which is the right order for a surface in a 64^3 grid. Sampler details worth recording, both from upstream: - CFG is a LERP (g*pos + (1-g)*neg), NOT neg + g*(pos-neg). Those differ non-linearly in strength rather than failing outright, so it is silent when wrong. Tested. - Guidance interval forces strength to 1 outside its window, which halves model calls there: 6 calls for 4 steps rather than 8. Tested by counting. Schedule matches upstream to 1e-16 and constant velocity integrates exactly at any step count. Remaining for a real image->3D run is CONDITIONING, not models: DINOv3 features plus Pixal3D's camera back-projection for the view-aligned 'proj' half. Deliberately not porting DINOv3 - it is a stock ViT run once per image, outside the denoising loop, so torch on MPS is the right tool and transformers gives exact parity for free.
119 lines
4.4 KiB
Python
119 lines
4.4 KiB
Python
"""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
|
|
|
|
# -- 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,
|
|
):
|
|
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))
|
|
return g * pos + (1 - g) * neg # LERP, matching upstream
|
|
|
|
# -- 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,
|
|
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
|
|
)
|
|
x = _like(x, _feats(x) - (t - t_prev) * v)
|
|
mx.eval(_feats(x))
|
|
if progress:
|
|
progress(i + 1, steps)
|
|
return x
|