SLAT stage: image -> occupancy -> sparse latents -> mesh, running end to end

The whole geometry chain now runs on a real photograph:

  [1] occupancy   12948 voxels                     19.7s
  [2] cond        proj (1, 262144, 2048) @ 64^3     2.6s
      gathered    proj (12948, 2048)
  [3] SLAT        (12948, 32)                      88.7s
  [4] MESH        3556515 verts, 7071196 faces     11.0s   grid 1024^3
      peak 22.6 GB, bounds inside the unit cube

Two bugs fixed on the way:

1. "global" must be FLAT [M,C] for the sparse blocks, not the dense stage's [B,T,C].
   The sparse cross-attention takes a token stack plus an explicit layout, so the
   dense shape dies inside to_kv's reshape rather than anywhere informative. Gathering
   now reshapes it, and refuses batch > 1 rather than silently mislabelling a layout.
2. o_voxel needs the decoder OUTPUT grid, not its configured resolution. The shape
   decoder applies four 2x upsamples, so a res-64 latent decodes into 1024^3, while
   the config says 256 (upstream overrides it per run via set_resolution). Passing 256
   raised an opaque out-of-bounds inside o_voxel's hashmap insert. Added
   output_resolution() and a guard that names the real cause.

HONEST LIMITATION - this is NOT yet the shipped cascade. Upstream's
sample_shape_slat_cascade runs the 512 flow (res 32) first, denormalises, UPSAMPLES
THE COORDINATE SET through the shape decoder, then runs the 1024 flow on the refined
coords. Running the HR flow straight off the 64^3 occupancy set yields a complete,
exportable mesh whose silhouette IoU is 0.639 - against 0.842 for the occupancy grid
that seeded it. The gap is a halo of geometry outside the true silhouette, exactly
what the missing coordinate refinement would prune. Do not read the current mesh
quality as the model's; wiring the cascade is the next step.

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

.gitignore

@@ -4,3 +4,9 @@ ckpt/
 weights/
 __pycache__/
 *.pyc
+pixal3d_geometry.glb
+mesh_silhouette.png
+silhouette_check.png
+pixal3d_geometry.glb
+mesh_silhouette.png
+silhouette_check.png

pixal3d_mlx/mesh.py

@@ -34,14 +34,35 @@ def _torch(a):
     return torch.from_numpy(np.asarray(a))
 
 
+def output_resolution(h, upsample_factor: int = 16) -> int:
+    """The decoder's OUTPUT grid size, which is what o_voxel needs.
+
+    The shape decoder applies four 2x upsamples, so a resolution-64 latent decodes into
+    a 1024^3 grid. The `resolution` field in the checkpoint config is the decoder's
+    configured default (256) — upstream overrides it per run via `set_resolution`, so
+    reading it off the config gives the wrong grid and o_voxel's hashmap then raises an
+    opaque out-of-bounds deep inside `insert`.
+    """
+    return int(mx.max(h.coords[:, 1:]).item()) // upsample_factor * upsample_factor + upsample_factor
+
+
 def fdg_to_mesh(h, resolution: int, voxel_margin: float = 0.5) -> Tuple:
     """Shape-decoder output -> (vertices, faces) as torch tensors.
 
     `h` is the decoder's SparseTensor: `h.feats` [N,7], `h.coords` [N,4] with the batch
-    index in column 0. Single batch item only, which is all inference ever produces.
+    index in column 0. `resolution` is the OUTPUT grid size (see `output_resolution`),
+    not the decoder's configured one. Single batch item only, which is all inference
+    ever produces.
     """
     from o_voxel.convert import flexible_dual_grid_to_mesh
 
+    hi = int(mx.max(h.coords[:, 1:]).item())
+    if hi >= resolution:
+        raise ValueError(
+            f"coords reach {hi} but grid_size={resolution}; pass the decoder's OUTPUT "
+            f"resolution (input_res * 16), not its configured default"
+        )
+
     feats = h.feats
     m = voxel_margin
     vertices = (1 + 2 * m) * mx.sigmoid(feats[..., 0:3]) - m

pixal3d_mlx/pipeline.py

@@ -26,6 +26,7 @@ from __future__ import annotations
 from typing import Callable, Optional
 
 import mlx.core as mx
+import numpy as np
 
 from .sampler import FlowEulerSampler
 
@@ -85,6 +86,79 @@ def run_structure_stage(
     return occ, latent
 
 
+def gather_proj_at_coords(cond: dict, coords: mx.array, grid_resolution: int) -> dict:
+    """Reduce the dense [B, R^3, C] proj grid to just the occupied voxels.
+
+    The structure stage conditions on every voxel of the lattice; the SLAT stage is
+    sparse and only ever touches the occupied set, so the proj features are gathered
+    at those coordinates. Coordinate order is (batch, x, y, z), matching upstream's
+    `encode_image_proj`.
+
+    `global` is not indexed — it is cross-attended — but it IS reshaped, because the
+    sparse attention path takes a flat [M, C] stack with an explicit layout, whereas
+    the dense structure stage takes [B, T, C]. Handing the dense shape to the sparse
+    blocks fails inside `to_kv`'s reshape rather than anywhere informative.
+    """
+    r = grid_resolution
+    proj = cond["proj"]
+    b_size = proj.shape[0]
+    if b_size != 1:
+        raise NotImplementedError(
+            "sparse conditioning assumes batch 1; a real batch needs an explicit "
+            "context layout rather than the default single slice"
+        )
+    proj = proj.reshape(b_size, r, r, r, -1)
+    b, x, y, z = (coords[:, i] for i in range(4))
+    flat = ((b * r + x) * r + y) * r + z
+    gathered = proj.reshape(-1, proj.shape[-1])[flat]
+    g = cond["global"]
+    return {"global": g.reshape(-1, g.shape[-1]), "proj": gathered}
+
+
+def run_slat_stage(
+    flow_model,
+    cond: dict,
+    coords: mx.array,
+    neg_cond=None,
+    normalization: dict | None = None,
+    seed: int = 0,
+    progress: Optional[Callable[[int, int], None]] = None,
+    **overrides,
+):
+    """Sparse latent sampling seeded by the occupancy set. Returns a SparseTensor.
+
+    The denormalisation at the end is NOT cosmetic: the flows are trained on
+    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.
+    """
+    from trellis_sparse_mlx import SparseTensor
+
+    params = {**SHAPE_SLAT_PARAMS, **overrides}
+    mx.random.seed(seed)
+    noise = SparseTensor(
+        mx.random.normal((coords.shape[0], flow_model.in_channels)), coords
+    )
+    slat = FlowEulerSampler().sample(
+        flow_model, noise, cond=cond, neg_cond=neg_cond, progress=progress, **params
+    )
+    if normalization:
+        mean = mx.array(np.asarray(normalization["mean"], dtype=np.float32))[None]
+        std = mx.array(np.asarray(normalization["std"], dtype=np.float32))[None]
+        slat = slat.replace(slat.feats * std + mean)
+    mx.eval(slat.feats)
+    return slat
+
+
 def image_to_occupancy(
     image_path: str,
     flow_model,