The shipped cascade: image -> GLB at silhouette IoU 0.969

image_to_mesh() now runs the real cascade, not the single-stage shortcut:

  structure    3048 voxels @32^3      (64^3 occupancy, MAX-POOLED DOWN)
  LR SLAT      3048 x 32              shape_512 extractor
  refine      13147 coords @64^3      four decoder stages -> coords -> quantise
  HR SLAT     13147 x 32              shape_1024 extractor
  mesh      3988052 verts, 7996876 faces @1024^3

  TOTAL 258.4s, peak 27.9GB with every model resident   silhouette IoU 0.969

Three things the cascade needed:

1. occupied_coords_at() - ss_dec always decodes 64^3 but the cascade STARTS at 32^3.
   Upstream max-pools the boolean grid down by the ratio (a voxel survives if ANY of
   its eight children was occupied). I had been feeding the raw 64^3 set to the HR flow.
2. decoder.upsample() - pushes the LR latent four stages in and returns COORDS, not
   features. The predicted subdivisions grow the occupied set; those coords quantise
   onto the HR flow's grid. Stops BEFORE stage `upsample_times`, as upstream does;
   one stage further doubles the resolution and misplaces every voxel.
3. grid_resolution override on ProjConditioner - upstream backs the HR grid off in
   128-unit steps while the token count exceeds max_num_tokens, so a dense object
   degrades instead of exploding. refine_coords() implements that loop.

I WAS WRONG ABOUT THE HALO. The previous commit blamed the single-stage shortcut for a
0.639 silhouette IoU and predicted the cascade would fix it. The cascade measured
0.640 - no change. The real fault was in my VERIFICATION, not the pipeline: o_voxel
returns vertices in the voxel-grid frame, while ProjGrid rotates its lattice by
_BLENDER_ROT before projecting. Rotating the mesh the same way scores 0.969 on the
same geometry the earlier commit had already produced. Added mesh.to_camera_frame()
so the trap is named where it bites; the earlier mesh was correct all along.

The cascade is still the right thing - it is the shipped path, and staged loading
halves peak memory (12.8GB vs 22.6GB) when models are released between stages.

Also adds models.load_all(), so a server builds all five models plus both conditioners
ONCE. Warmup is ~71s against ~17s of compute, so an operator must never fork per job.
Holding everything resident costs 27.9GB peak - nothing on a 256GB box.

scripts/image_to_mesh.py exits non-zero if IoU < 0.85: a run that completes with a bad
reconstruction has failed even though nothing raised.

27/27 green.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

.gitignore

@@ -10,3 +10,6 @@ silhouette_check.png
 pixal3d_geometry.glb
 mesh_silhouette.png
 silhouette_check.png
+output.glb
+pixal3d_cascade.glb
+cascade_silhouette.png

pixal3d_mlx/cond.py

@@ -98,14 +98,29 @@ class ProjConditioner:
                                   image_resolution=self.image_size)
         self.proj_channels = self.encoder.embed_dim * (2 if self.use_naf else 1)
 
+    def _grid(self, grid_resolution: int | None):
+        """ProjGrid at `grid_resolution`, cached. Always at image_size (see __init__)."""
+        r = grid_resolution or self.grid_resolution
+        if r == self.grid_resolution:
+            return self.proj_grid
+        cache = getattr(self, "_grid_cache", None) or {}
+        if r not in cache:
+            cache[r] = ProjGrid(grid_resolution=r, image_resolution=self.image_size)
+            self._grid_cache = cache
+        return cache[r]
+
     def __call__(self, image_chw01: np.ndarray, camera_angle_x: float,
-                 distance: float | None = None, mesh_scale: float = 1.0
+                 distance: float | None = None, mesh_scale: float = 1.0,
+                 grid_resolution: int | None = None,
                  ) -> Tuple[Dict[str, mx.array], Dict[str, mx.array]]:
+        """`grid_resolution` overrides this stage's lattice — the cascade backs the
+        high-res grid off in steps when an object is too dense for max_num_tokens."""
         if distance is None:
             distance = distance_from_fov(camera_angle_x, mesh_scale, self.image_size)
 
+        grid = self._grid(grid_resolution)
         z_global, z_patch = self.encoder(image_chw01)
-        z_proj = self.proj_grid(z_patch, camera_angle_x, distance, mesh_scale)
+        z_proj = grid(z_patch, camera_angle_x, distance, mesh_scale)
 
         if self.use_naf:
             from .naf import upsample
@@ -122,7 +137,7 @@ class ProjConditioner:
 
             hr = upsample(z_patch, guide, self.naf_target,
                           device=self.device or "mps")
-            z_hr = self.proj_grid(hr, camera_angle_x, distance, mesh_scale)
+            z_hr = grid(hr, camera_angle_x, distance, mesh_scale)
             z_proj = mx.concatenate([z_proj, z_hr], axis=-1)
 
         cond = {"global": z_global, "proj": z_proj}

pixal3d_mlx/decoders.py

@@ -114,6 +114,31 @@ class SparseUnetVaeDecoder(nn.Module):
         out = self.output_layer(h.replace(f))
         return (out, subs) if return_subs else out
 
+    def upsample(self, x: SparseTensor, upsample_times: int) -> mx.array:
+        """Run only the first `upsample_times` stages and return the COORDS.
+
+        The cascade needs a refined coordinate set, not features: the low-res SLAT is
+        pushed part-way through the decoder purely so its predicted subdivisions grow
+        the occupied set, and the resulting coords seed the high-res flow. Each stage
+        doubles resolution, so `upsample_times=4` yields coords at 16x the input.
+
+        Deliberately stops BEFORE running stage `upsample_times`, matching upstream —
+        running one stage further would return coords at 2x the intended resolution
+        and silently misplace every voxel in the high-res pass.
+        """
+        if not self.pred_subdiv:
+            raise ValueError("upsample needs a decoder with pred_subdiv=True")
+        h = self.from_latent(x)
+        for i, stage in enumerate(self.blocks):
+            if i == upsample_times:
+                return h.coords
+            for j, blk in enumerate(stage):
+                if i < len(self.blocks) - 1 and j == len(stage) - 1:
+                    h, _ = blk(h)
+                else:
+                    h = blk(h)
+        return h.coords
+
 
 class FlexiDualGridVaeDecoder(SparseUnetVaeDecoder):
     """shape_dec. Emits 7 channels; the first 3 are vertex offsets within the voxel."""

pixal3d_mlx/mesh.py

@@ -115,3 +115,17 @@ EXPORT_ROTATION = np.array([[-1, 0, 0, 0],
                             [0, 0, -1, 0],
                             [0, -1, 0, 0],
                             [0, 0, 0, 1]], dtype=np.float64)
+
+
+def to_camera_frame(vertices):
+    """Mesh vertices -> the frame `proj.project_points` expects.
+
+    THE GOTCHA: o_voxel returns vertices in the VOXEL GRID's frame — a linear map from
+    integer coords into the aabb. `ProjGrid` rotates its lattice by `_BLENDER_ROT`
+    BEFORE projecting, so mesh vertices must be rotated the same way to be compared
+    against the source image. Skipping this does not throw; it silently reprojects a
+    rotated object, which reads as a plausible-looking blob with a halo. It cost a
+    wrong diagnosis here: a correct 0.969 silhouette IoU measured as 0.640.
+    """
+    from .proj import _BLENDER_ROT
+    return np.asarray(vertices) @ _BLENDER_ROT.T

pixal3d_mlx/models.py

@@ -0,0 +1,68 @@
+"""Load every model once.
+
+Warmup dominates a single run — roughly 71s of graph build and weight fault against
+~17s of actual compute for the structure stage alone. So anything serving more than
+one job (the MODELBEAST operator, a batch script) must build this ONCE and keep it,
+never fork per job. The trellis-2 lane on this fleet shows the same shape: 47.9s cold
+against 2.5s warm pipeline load.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+CKPTS = Path("/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts")
+PIPELINE_JSON = CKPTS.parent / "pipeline.json"
+REPO_WEIGHTS = Path(__file__).resolve().parents[1] / "weights"
+
+FILES = {
+    "ss_flow": "ss_flow_img_dit_1_3B_64_bf16",
+    "ss_dec": "ss_dec_conv3d_16l8_fp16",
+    "slat_512": "slat_flow_img2shape_dit_1_3B_512_bf16",
+    "slat_1024": "slat_flow_img2shape_dit_1_3B_1024_bf16",
+    "shape_dec": "shape_dec_next_dc_f16c32_fp16",
+    "tex_dec": "tex_dec_next_dc_f16c32_fp16",
+}
+
+
+def _weights(stem: str) -> Path:
+    """Decoders were converted into the repo; flows pass through untouched."""
+    local = REPO_WEIGHTS / f"{stem}.safetensors"
+    return local if local.exists() else CKPTS / f"{stem}.safetensors"
+
+
+def normalization(kind: str = "shape") -> dict:
+    cfg = json.loads(PIPELINE_JSON.read_text())
+    return cfg.get("args", cfg)[f"{kind}_slat_normalization"]
+
+
+def load_all(device: str | None = None, with_texture: bool = False) -> dict:
+    """Every model plus both conditioners. ~24GB of weights; MLX loads them lazily."""
+    from . import slat_flow, ss_dec, ss_flow
+    from .cond import ProjConditioner
+    from .decoders import load as load_dec
+
+    def cfg(stem):
+        return CKPTS / f"{stem}.json"
+
+    models = {}
+    for key, loader in (("ss_flow", ss_flow.load), ("ss_dec", ss_dec.load),
+                        ("slat_512", slat_flow.load), ("slat_1024", slat_flow.load)):
+        stem = FILES[key]
+        model, rep = loader(_weights(stem), cfg(stem))
+        if rep["missing"] or rep["unmapped"]:
+            raise RuntimeError(f"{key}: {len(rep['missing'])} missing, "
+                               f"{len(rep['unmapped'])} unmapped")
+        models[key] = model
+
+    for key in (["shape_dec", "tex_dec"] if with_texture else ["shape_dec"]):
+        stem = FILES[key]
+        model, rep = load_dec(_weights(stem), cfg(stem))
+        if rep["missing"] or rep["unmapped"]:
+            raise RuntimeError(f"{key}: incomplete weight mapping")
+        models[key] = model
+
+    models["cond_512"] = ProjConditioner("shape_512", device=device)
+    models["cond_1024"] = ProjConditioner("shape_1024", device=device)
+    return models

pixal3d_mlx/pipeline.py

@@ -23,6 +23,7 @@ upstream itself falls back to the ungated `camenduru` mirror.
 
 from __future__ import annotations
 
+import time
 from typing import Callable, Optional
 
 import mlx.core as mx
@@ -131,15 +132,8 @@ def run_slat_stage(
     standardised latents, and the decoders expect raw ones. Skipping it yields a mesh
     that decodes without error and is quietly the wrong scale.
 
-    SINGLE-STAGE ONLY — this is not yet the cascade upstream actually ships. Its
-    `sample_shape_slat_cascade` runs the 512 flow (resolution 32) first, denormalises,
-    UPSAMPLES THE COORDINATE SET through the shape decoder, and only then runs the 1024
-    flow (resolution 64) on the refined coords. Running the HR flow straight off the
-    64^3 occupancy set produces a complete, exportable mesh whose silhouette IoU is
-    0.639 against 0.842 for the occupancy grid that seeded it — the loss is a halo of
-    geometry outside the true silhouette, which is what the missing coordinate
-    refinement would have pruned. Wiring the cascade is the next step; do not read the
-    current mesh quality as the model's.
+    One stage of the cascade. `image_to_mesh` composes two of these — see there for the
+    LR -> refine -> HR sequence upstream actually ships.
     """
     from trellis_sparse_mlx import SparseTensor
 
@@ -159,6 +153,116 @@ def run_slat_stage(
     return slat
 
 
+def refine_coords(shape_decoder, lr_slat, hr_resolution: int = 1024,
+                  lr_resolution: int = 512, max_num_tokens: int = 49152):
+    """Low-res SLAT -> refined high-res coordinate set. Stage 3a of the cascade.
+
+    The low-res latent is pushed four stages into the shape decoder purely so its
+    predicted subdivisions grow the occupied set; those coords (now at 16x, i.e. 512)
+    are then quantised down to the high-res flow's own grid (1024 // 16 = 64) and
+    deduplicated. This is the step whose absence leaves geometry outside the silhouette.
+
+    Upstream backs the resolution off in 128 steps while the token count exceeds
+    `max_num_tokens`, so a dense object degrades gracefully instead of exploding.
+    Returns (coords, hr_resolution, grid_res).
+    """
+    import numpy as _np
+
+    hr_coords = _np.asarray(shape_decoder.upsample(lr_slat, upsample_times=4))
+    res = hr_resolution
+    while True:
+        grid_res = res // 16
+        q = _np.concatenate([
+            hr_coords[:, :1],
+            _np.rint((hr_coords[:, 1:] + 0.5) / lr_resolution * (grid_res - 1)).astype(_np.int32),
+        ], axis=1)
+        uniq = _np.unique(q, axis=0)
+        if uniq.shape[0] < max_num_tokens or res == 1024:
+            break
+        res -= 128
+    return mx.array(uniq.astype(_np.int32)), res, res // 16
+
+
+def image_to_mesh(
+    image_path: str,
+    models: dict,
+    camera_angle_x: float = DEFAULT_FOV,
+    mesh_scale: float = 1.0,
+    seed: int = 0,
+    normalization: dict | None = None,
+    log=print,
+):
+    """The full shipped geometry cascade: image -> (vertices, faces).
+
+    `models` supplies the five pieces, so a long-lived server can load them ONCE:
+    `ss_flow`, `ss_dec`, `slat_512`, `slat_1024`, `shape_dec`, plus the conditioners
+    `cond_512` and `cond_1024`.
+
+    The sequence, and why each step is where it is:
+
+      1. structure  -> 64^3 occupancy, MAX-POOLED DOWN to 32^3. The cascade starts low.
+      2. LR SLAT    -> sparse latents at 32^3, conditioned by the shape_512 extractor
+      3. refine     -> the LR latent is pushed four decoder stages so its predicted
+                       subdivisions grow the occupied set; those coords quantise to 64^3
+      4. HR SLAT    -> sparse latents at the refined coords, shape_1024 extractor
+      5. decode     -> Flexible Dual Grid at 1024^3 -> triangles
+
+    Returns (vertices, faces, info). Vertices are in the VOXEL GRID frame; use
+    `mesh.to_camera_frame` before comparing against the source image.
+    """
+    from .cond import load_image
+    from .mesh import fdg_to_mesh, output_resolution
+
+    t0 = time.time()
+    occ, _, _ = image_to_occupancy(image_path, models["ss_flow"], models["ss_dec"],
+                                   camera_angle_x=camera_angle_x,
+                                   mesh_scale=mesh_scale, seed=seed)
+    coords = occupied_coords_at(occ, 32)
+    log(f"  structure  {coords.shape[0]:6} voxels @32^3   {time.time() - t0:6.1f}s")
+
+    t = time.time()
+    c512 = models["cond_512"]
+    img = load_image(image_path, c512.image_size)
+    cond, uncond = c512(img, camera_angle_x, mesh_scale=mesh_scale)
+    lr_slat = run_slat_stage(
+        models["slat_512"], gather_proj_at_coords(cond, coords, 32), coords,
+        neg_cond=gather_proj_at_coords(uncond, coords, 32),
+        normalization=normalization, seed=seed)
+    log(f"  LR SLAT    {lr_slat.feats.shape[0]:6} x {lr_slat.feats.shape[1]}          "
+        f"{time.time() - t:6.1f}s")
+
+    t = time.time()
+    hr_coords, hr_res, grid_res = refine_coords(models["shape_dec"], lr_slat)
+    log(f"  refine     {hr_coords.shape[0]:6} coords @{grid_res}^3 (res {hr_res})  "
+        f"{time.time() - t:6.1f}s")
+
+    t = time.time()
+    c1024 = models["cond_1024"]
+    img = load_image(image_path, c1024.image_size)
+    cond, uncond = c1024(img, camera_angle_x, mesh_scale=mesh_scale,
+                         grid_resolution=grid_res)
+    hr_slat = run_slat_stage(
+        models["slat_1024"], gather_proj_at_coords(cond, hr_coords, grid_res), hr_coords,
+        neg_cond=gather_proj_at_coords(uncond, hr_coords, grid_res),
+        normalization=normalization, seed=seed)
+    log(f"  HR SLAT    {hr_slat.feats.shape[0]:6} x {hr_slat.feats.shape[1]}          "
+        f"{time.time() - t:6.1f}s")
+
+    t = time.time()
+    out, subs = models["shape_dec"](hr_slat, return_subs=True)
+    mx.eval(out.feats)
+    res_out = output_resolution(out)
+    v, f = fdg_to_mesh(out, res_out)
+    log(f"  mesh       {v.shape[0]:6} verts, {f.shape[0]} faces @{res_out}^3  "
+        f"{time.time() - t:6.1f}s")
+
+    info = {"hr_resolution": hr_res, "grid_resolution": grid_res,
+            "output_resolution": res_out, "num_tokens": int(hr_coords.shape[0]),
+            "seconds": round(time.time() - t0, 1),
+            "peak_gb": round(mx.get_peak_memory() / 2 ** 30, 1)}
+    return v, f, {**info, "subs": subs, "hr_slat": hr_slat}
+
+
 def image_to_occupancy(
     image_path: str,
     flow_model,
@@ -195,6 +299,28 @@ def image_to_occupancy(
     return occ, latent, cond
 
 
+def occupied_coords_at(occ: mx.array, resolution: int, threshold: float = 0.0) -> mx.array:
+    """Occupancy logits -> int32 [N,4] (batch, x, y, z) coords at `resolution`.
+
+    ss_dec always decodes to 64^3, but the cascade STARTS AT 32^3 — upstream max-pools
+    the boolean grid down by the ratio (`max_pool3d(..., ratio) > 0.5`, i.e. a voxel
+    survives if ANY of its eight children was occupied). Feeding the raw 64^3 set into
+    the high-res flow instead skips the low-res refinement entirely and leaves a halo
+    of geometry outside the true silhouette.
+    """
+    import numpy as _np
+
+    o = _np.asarray(occ)[:, 0] > threshold          # [B,D,H,W] bool
+    d = o.shape[1]
+    if resolution != d:
+        ratio = d // resolution
+        if d % resolution:
+            raise ValueError(f"{d} is not an integer multiple of {resolution}")
+        o = o.reshape(o.shape[0], resolution, ratio, resolution, ratio, resolution, ratio)
+        o = o.any(axis=(2, 4, 6))                   # max-pool over the boolean grid
+    return mx.array(_np.argwhere(o).astype(_np.int32))
+
+
 def occupied_coords(occ: mx.array, threshold: float = 0.0) -> mx.array:
     """Occupancy logits -> int32 [N, 4] (batch, z, y, x) coords for the SLAT stage."""
     import numpy as np

scripts/image_to_mesh.py

@@ -0,0 +1,92 @@
+"""Image -> GLB through the full Pixal3D cascade, with a silhouette check.
+
+Usage:  python scripts/image_to_mesh.py IMAGE [-o OUT.glb] [--fov RAD] [--seed N]
+
+The silhouette IoU is the acceptance test: re-project the mesh through the same camera
+and compare against the input matte. Pixal3D's entire claim is pixel alignment, so a
+run that completes with a poor IoU has failed even though nothing raised.
+"""
+
+import argparse
+import sys
+import time
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+import trimesh
+
+REPO = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO))
+
+from pixal3d_mlx.mesh import to_camera_frame                       # noqa: E402
+from pixal3d_mlx.models import load_all, normalization             # noqa: E402
+from pixal3d_mlx.pipeline import DEFAULT_FOV, image_to_mesh        # noqa: E402
+
+DEFAULT_IMAGE = REPO / "upstream" / "Pixal3D" / "assets" / "images" / "0_img.png"
+
+
+def silhouette_iou(vertices, image_path, fov, res=512):
+    """Re-project the mesh through the generating camera; IoU against the input matte.
+
+    Vertices MUST be rotated into the camera frame first — o_voxel returns them in the
+    voxel-grid frame while ProjGrid rotates its lattice before projecting.
+    """
+    from PIL import Image
+    from scipy.ndimage import binary_dilation
+
+    from pixal3d_mlx.cond import preprocess_image
+    from pixal3d_mlx.proj import _FRONT_VIEW, distance_from_fov, project_points
+
+    img = Image.open(image_path)
+    if img.mode != "RGBA":
+        return None
+    matte = np.asarray(preprocess_image(img).convert("L").resize((res, res))) > 8
+
+    tm = _FRONT_VIEW.copy()
+    tm[1, 3] = -distance_from_fov(fov, 1.0, res)
+    pts = to_camera_frame(vertices).astype(np.float32)[None]
+    px, _, _ = project_points(mx.array(pts), mx.array(tm[None]), fov, res)
+    px = np.asarray(px)[0].astype(int)
+
+    keep = (px[:, 0] >= 0) & (px[:, 0] < res) & (px[:, 1] >= 0) & (px[:, 1] < res)
+    proj = np.zeros((res, res), bool)
+    proj[px[keep, 1], px[keep, 0]] = True
+    proj = binary_dilation(proj, np.ones((3, 3), bool))
+    return (proj & matte).sum() / (proj | matte).sum(), proj, matte
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("image", nargs="?", default=str(DEFAULT_IMAGE))
+    ap.add_argument("-o", "--output", default=str(REPO / "output.glb"))
+    ap.add_argument("--fov", type=float, default=DEFAULT_FOV)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--no-check", action="store_true")
+    a = ap.parse_args()
+
+    t = time.time()
+    models = load_all()
+    print(f"loaded models ({time.time() - t:.1f}s, lazy — weights fault in on first use)")
+
+    v, f, info = image_to_mesh(a.image, models, camera_angle_x=a.fov, seed=a.seed,
+                               normalization=normalization("shape"))
+
+    mesh = trimesh.Trimesh(v.cpu().numpy(), f.cpu().numpy(), process=False)
+    mesh.export(a.output)
+    print(f"\nTOTAL {info['seconds']}s   peak {info['peak_gb']} GB   -> {a.output}")
+
+    if not a.no_check:
+        got = silhouette_iou(mesh.vertices, a.image, a.fov)
+        if got is None:
+            print("(input has no alpha matte — skipping silhouette check)")
+        else:
+            iou, _, _ = got
+            print(f"silhouette IoU {iou:.3f}")
+            if iou < 0.85:
+                print("WARNING: low IoU — the reconstruction is not tracking the input")
+                sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()