diff --git a/.gitignore b/.gitignore index a379c42..dd9e07c 100644 --- a/.gitignore +++ b/.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 diff --git a/pixal3d_mlx/cond.py b/pixal3d_mlx/cond.py index d7d0d06..d061f28 100644 --- a/pixal3d_mlx/cond.py +++ b/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} diff --git a/pixal3d_mlx/decoders.py b/pixal3d_mlx/decoders.py index a11acd3..7d3dd58 100644 --- a/pixal3d_mlx/decoders.py +++ b/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.""" diff --git a/pixal3d_mlx/mesh.py b/pixal3d_mlx/mesh.py index cb321ca..52d653a 100644 --- a/pixal3d_mlx/mesh.py +++ b/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 diff --git a/pixal3d_mlx/models.py b/pixal3d_mlx/models.py new file mode 100644 index 0000000..5936c9c --- /dev/null +++ b/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 diff --git a/pixal3d_mlx/pipeline.py b/pixal3d_mlx/pipeline.py index 6cc7854..d6420ff 100644 --- a/pixal3d_mlx/pipeline.py +++ b/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 diff --git a/scripts/image_to_mesh.py b/scripts/image_to_mesh.py new file mode 100644 index 0000000..80f263a --- /dev/null +++ b/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()