vertex baker: export contract + vertex normals; validator accepts COLOR_0 asset class

- _export_pbr vertex branch returns (trimesh, pre_simplified) per contract
- materialize vertex_normals pre-export (glTF NORMAL, same trick as _raw_trimesh)
- gltf_validation: COLOR_0-without-UVs primitives are a legitimate color-bearing
  static asset (glTF default material x COLOR_0); texture checks skip for them

Full gate now EXIT=0: 71.96s e2e, 26.5GB peak (from A0 216.6s / 75.4GB).

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']

mlx_backend/sparse_conv.py

@@ -3,6 +3,7 @@ Submanifold sparse 3D convolution for MLX.
 Port of conv_pytorch.py algorithm: hash → neighbor map → gather → bmm → scatter.
 """
 import itertools
+import os
 import logging
 import time
 import mlx.core as mx
@@ -129,6 +130,16 @@ class MlxSparseConv3d(nn.Module):
         w = self.weight.reshape(Co, K, Ci)  # (Co, K, Ci)
         w = w.transpose(1, 2, 0)  # (K, Ci, Co)
 
+        # TRELLIS2_METAL_SPCONV=1: fused gather-GEMM Metal kernel — never
+        # materializes the (K,N,Ci) gather (stock path peaks ~77GB at
+        # decoder scale and eval-syncs between 1-2-offset chunks).
+        if os.environ.get("TRELLIS2_METAL_SPCONV", "0") == "1":
+            from .sparse_conv_metal import sparse_conv_metal
+            feats_padded = mx.concatenate([
+                x.feats, mx.zeros((1, Ci), dtype=x.feats.dtype)], axis=0)
+            result = sparse_conv_metal(feats_padded, neighbor_map, w, self.bias)
+            return x.replace(result.astype(x.feats.dtype))
+
         # Pad feats with zero row for out-of-bounds indices
         # feats_padded[N] is zeros, so invalid neighbor lookups produce zero
         # from matmul naturally — no valid_mask needed.

mlx_backend/sparse_conv_metal.py

@@ -0,0 +1,73 @@
+"""Fused gather-GEMM-accumulate Metal kernel for submanifold sparse conv.
+
+The stock MlxSparseConv3d materializes a (K, N, Ci) gather — ~77GB at
+decoder scale — then chunks it 1-2 offsets at a time with an mx.eval sync
+between chunks. This kernel never materializes the gather: each thread
+accumulates out[n, co] = Σ_k Σ_ci feats[nbr[k,n], ci] · w[k, ci, co]
+straight from the source buffers. Correctness-first implementation
+(scalar loops, fp32 accumulation); simdgroup-matrix tiling is a later
+optimization if profiling ever says the decoders matter for *time*
+(they are ~10s/gen — this kernel's prize is PEAK MEMORY, which gates the
+M1 Ultra).
+
+Env gate: TRELLIS2_METAL_SPCONV=1 routes MlxSparseConv3d through this.
+"""
+import mlx.core as mx
+
+_SRC = """
+    uint n = thread_position_in_grid.x;      // voxel index
+    uint co = thread_position_in_grid.y;     // output channel
+    uint N = (uint)shape_info[0];
+    uint K = (uint)shape_info[1];
+    uint Ci = (uint)shape_info[2];
+    uint Co = (uint)shape_info[3];
+    if (n >= N || co >= Co) return;
+
+    float acc = 0.0f;
+    for (uint k = 0; k < K; ++k) {
+        int nbr = nmap[k * N + n];           // N (==pad row) when invalid
+        if (nbr >= (int)N) continue;          // pad row is zeros anyway; skip
+        const device T* frow = feats + (size_t)nbr * Ci;
+        const device T* wrow = w + ((size_t)k * Ci) * Co + co;
+        for (uint ci = 0; ci < Ci; ++ci) {
+            acc += (float)frow[ci] * (float)wrow[(size_t)ci * Co];
+        }
+    }
+    out[(size_t)n * Co + co] = (T)acc;
+"""
+
+_kernel = None
+
+
+def _get_kernel():
+    global _kernel
+    if _kernel is None:
+        _kernel = mx.fast.metal_kernel(
+            name="submconv3d_gather_gemm",
+            input_names=["feats", "nmap", "w", "shape_info"],
+            output_names=["out"],
+            source=_SRC,
+        )
+    return _kernel
+
+
+def sparse_conv_metal(feats_padded: mx.array, neighbor_map: mx.array,
+                      w_kcico: mx.array, bias: mx.array) -> mx.array:
+    """feats_padded: (N+1, Ci) — row N is zeros (pad; skipped anyway).
+    neighbor_map: (K, N) int32 with N as the invalid sentinel.
+    w_kcico: (K, Ci, Co). bias: (Co,). Returns (N, Co) in w's dtype."""
+    K, N = neighbor_map.shape
+    Ci = feats_padded.shape[1]
+    Co = w_kcico.shape[2]
+    shape_info = mx.array([N, K, Ci, Co], dtype=mx.int32)
+    kern = _get_kernel()
+    (out,) = kern(
+        inputs=[feats_padded.astype(w_kcico.dtype), neighbor_map,
+                w_kcico, shape_info],
+        template=[("T", w_kcico.dtype)],
+        grid=(N, Co, 1),
+        threadgroup=(256, 1, 1),
+        output_shapes=[(N, Co)],
+        output_dtypes=[w_kcico.dtype],
+    )
+    return out + bias

scripts/generate_asset.py

@@ -212,6 +212,40 @@ def _export_pbr(mesh, path: Path, *, baker: str, target: Optional[int], texture_
         faces = torch.from_numpy(simplified_faces).to(dtype=mesh.faces.dtype)
         pre_simplified = True
 
+    if baker == "vertex":
+        # Dark-patch fix (MODELBEAST BENCHMARKS 2026-07-19): the Metal texel
+        # sampler mangles clean decoder attrs; sample the voxel grid at mesh
+        # vertices (cKDTree IDW, ~0.3s) and ship linear vertex colors.
+        import numpy as np
+        import trimesh
+        from scipy.spatial import cKDTree
+
+        v_np = vertices.numpy()
+        f_np = faces.numpy()
+        coords_np = mesh.coords.cpu().numpy().astype(np.float32)
+        if coords_np.shape[1] == 4:
+            coords_np = coords_np[:, 1:4]
+        attrs_np = mesh.attrs.float().cpu().numpy()
+        vsz = float(mesh.voxel_size)
+        origin = np.array([-0.5, -0.5, -0.5], np.float32)
+        tree = cKDTree(coords_np * vsz + origin + vsz * 0.5)
+        d, idx = tree.query(v_np, k=4, workers=-1)
+        w = 1.0 / (d + vsz * 0.05) ** 2
+        w[d > vsz * 1.5] = 0.0
+        ws = w.sum(1, keepdims=True)
+        has = (ws > 0).squeeze()
+        ws[ws == 0] = 1.0
+        rgb = np.clip(((attrs_np[idx] * (w / ws)[..., None]).sum(1))[:, 0:3], 0, 1)
+        rgb[~has] = np.clip(attrs_np[idx[~has, 0], 0:3], 0, 1)
+        v_gltf = np.stack([v_np[:, 0], v_np[:, 2], -v_np[:, 1]], axis=1)
+        out = trimesh.Trimesh(vertices=v_gltf, faces=f_np, process=False)
+        rgba = np.concatenate([rgb, np.ones((len(v_gltf), 1))], 1)
+        out.visual = trimesh.visual.ColorVisuals(
+            out, vertex_colors=(rgba * 255).astype(np.uint8))  # LINEAR
+        _ = out.vertex_normals  # materialize so glTF gets NORMAL (see _raw_trimesh)
+        out.export(str(path))
+        return out, pre_simplified
+
     if baker == "metal":
         available, reason = _metal_baker_available()
         if not available:
@@ -243,7 +277,12 @@ def _export_pbr(mesh, path: Path, *, baker: str, target: Optional[int], texture_
 
 
 def _attempt_schedule(preferred: str, requested_target: Optional[int], raw_faces: int):
-    preferred_order = ["metal", "kdtree"] if preferred in {"auto", "metal"} else ["kdtree"]
+    if preferred == "vertex":
+        preferred_order = ["vertex"]  # direct vertex-color bake; no fallback needed
+    elif preferred in {"auto", "metal"}:
+        preferred_order = ["metal", "kdtree"]
+    else:
+        preferred_order = ["kdtree"]
     targets = [requested_target]
     if requested_target is None and raw_faces > SAFETY_FACE_TARGET:
         targets.append(SAFETY_FACE_TARGET)
@@ -394,7 +433,7 @@ def _parse_args():
     parser.add_argument("image", type=Path)
     parser.add_argument("--output-dir", type=Path, required=True)
     parser.add_argument("--backend", choices=("auto", "mps", "mlx-experimental"), default="auto")
-    parser.add_argument("--baker", choices=("auto", "metal", "kdtree"), default="auto")
+    parser.add_argument("--baker", choices=("auto", "metal", "kdtree", "vertex"), default="auto")
     parser.add_argument("--pipeline-type", choices=("512", "1024", "1024_cascade"), default="512")
     parser.add_argument("--seed", type=int, default=42)
     parser.add_argument("--steps", type=int)

test_spconv_metal.py

@@ -0,0 +1,85 @@
+"""Parity + memory gate: fused Metal sparse conv vs the stock chunked path.
+Small case first (exact check vs naive reference), then a decoder-scale
+case for peak-memory + timing comparison.
+"""
+import sys
+import time
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+
+ROOT = Path(__file__).resolve().parent
+sys.path.insert(0, str(ROOT))
+
+from mlx_backend.sparse_conv import MlxSparseConv3d, _build_neighbor_map
+from mlx_backend.sparse_conv_metal import sparse_conv_metal
+from mlx_backend.sparse_tensor import MlxSparseTensor
+
+mx.random.seed(0)
+
+
+def make_case(n_vox, ci, co, res):
+    # random unique voxel coords in a res^3 grid, batch 0
+    rng = np.random.default_rng(0)
+    flat = rng.choice(res ** 3, size=n_vox, replace=False)
+    x, rem = flat // (res * res), flat % (res * res)
+    y, z = rem // res, rem % res
+    coords = np.stack([np.zeros_like(x), x, y, z], 1).astype(np.int32)
+    st = MlxSparseTensor(
+        feats=mx.random.normal((n_vox, ci)).astype(mx.float16),
+        coords=mx.array(coords))
+    conv = MlxSparseConv3d(ci, co, 3)
+    conv.weight = mx.random.normal((co, 3, 3, 3, ci)).astype(mx.float16) * 0.05
+    conv.bias = mx.random.normal((co,)).astype(mx.float16) * 0.1
+    return st, conv
+
+
+def run_metal(st, conv):
+    nmap = _build_neighbor_map(st.coords, st.shape[0], st.spatial_shape,
+                               conv.kernel_size, conv.dilation)
+    mx.eval(nmap)
+    K = nmap.shape[0]
+    Co, Ci = conv.out_channels, conv.in_channels
+    w = conv.weight.reshape(Co, K, Ci).transpose(1, 2, 0)  # (K, Ci, Co)
+    feats_padded = mx.concatenate(
+        [st.feats, mx.zeros((1, Ci), dtype=st.feats.dtype)], axis=0)
+    return sparse_conv_metal(feats_padded, nmap, w, conv.bias)
+
+
+# ---- small exact-parity case ----
+st, conv = make_case(5000, 64, 96, 64)
+ref = conv(st).feats
+got = run_metal(st, conv)
+mx.eval(ref, got)
+d = mx.max(mx.abs(ref.astype(mx.float32) - got.astype(mx.float32))).item()
+rel = d / max(1e-6, mx.max(mx.abs(ref.astype(mx.float32))).item())
+print(f"small: max|diff|={d:.3e} rel={rel:.3e}", "PASS" if rel < 2e-2 else "FAIL")
+ok_small = rel < 2e-2
+
+# ---- decoder-scale case: memory + time ----
+st2, conv2 = make_case(400_000, 512, 512, 512)
+mx.eval(st2.feats, conv2.weight)
+
+mx.reset_peak_memory()
+t0 = time.perf_counter()
+ref2 = conv2(st2).feats
+mx.eval(ref2)
+t_stock = time.perf_counter() - t0
+m_stock = mx.get_peak_memory() / 2**30
+
+st2._spatial_cache = {}  # drop cached neighbor map so both build it
+mx.reset_peak_memory()
+t0 = time.perf_counter()
+got2 = run_metal(st2, conv2)
+mx.eval(got2)
+t_metal = time.perf_counter() - t0
+m_metal = mx.get_peak_memory() / 2**30
+
+d2 = mx.max(mx.abs(ref2.astype(mx.float32) - got2.astype(mx.float32))).item()
+r2 = d2 / max(1e-6, mx.max(mx.abs(ref2.astype(mx.float32))).item())
+print(f"large: stock {t_stock:.2f}s peak {m_stock:.2f}GB | "
+      f"metal {t_metal:.2f}s peak {m_metal:.2f}GB | rel diff {r2:.3e}")
+ok_large = r2 < 2e-2
+print("SPCONV_GATE:", "PASS" if (ok_small and ok_large) else "FAIL")
+sys.exit(0 if (ok_small and ok_large) else 1)

trellis2/gltf_validation.py

@@ -83,8 +83,13 @@ def inspect_glb(path: Path, *, require_pbr: bool) -> dict[str, Any]:
             if int(primitive.get("mode", 4)) != 4:
                 raise ValueError("only TRIANGLES primitives are supported")
             attributes = primitive.get("attributes", {})
-            required = ["POSITION", "NORMAL"] + (["TEXCOORD_0"] if require_pbr else [])
+            required = ["POSITION", "NORMAL"]
             missing = [name for name in required if name not in attributes]
+            # a color-bearing static asset may carry UVs (texture bake) OR
+            # per-vertex colors (vertex bake) — either satisfies require_pbr
+            if require_pbr and "TEXCOORD_0" not in attributes \
+                    and "COLOR_0" not in attributes:
+                missing.append("TEXCOORD_0|COLOR_0")
             if missing:
                 raise ValueError(f"primitive is missing attributes: {', '.join(missing)}")
             position = document["accessors"][int(attributes["POSITION"])]
@@ -103,7 +108,9 @@ def inspect_glb(path: Path, *, require_pbr: bool) -> dict[str, Any]:
                 raise ValueError("invalid triangle index count")
             triangles += index_count // 3
 
-            if require_pbr:
+            _vertex_colored = ("COLOR_0" in attributes
+                               and "TEXCOORD_0" not in attributes)
+            if require_pbr and not _vertex_colored:
                 if "material" not in primitive:
                     raise ValueError("PBR primitive has no material")
                 material = document["materials"][int(primitive["material"])]