165de26db5
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>
69 lines
2.6 KiB
Python
69 lines
2.6 KiB
Python
"""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
|