diff --git a/.gitignore b/.gitignore index 4b3d9be..a379c42 100644 --- a/.gitignore +++ b/.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 diff --git a/pixal3d_mlx/mesh.py b/pixal3d_mlx/mesh.py index e331b6e..cb321ca 100644 --- a/pixal3d_mlx/mesh.py +++ b/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 diff --git a/pixal3d_mlx/pipeline.py b/pixal3d_mlx/pipeline.py index b9cc9d9..6cc7854 100644 --- a/pixal3d_mlx/pipeline.py +++ b/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,