b6f0d76619
The texture flow is imgshape2tex - it denoises 32 PBR channels while SEEING the shape latent, so in_channels is 64 against out_channels 32. Upstream feeds the shape latent as concat_cond and the model does sparse_cat([x, concat_cond], dim=-1); both share coords, so it reduces to a channel concat. Added to slat_flow and carried on the sampler (it is fixed for the whole trajectory and must reach BOTH CFG branches). Verified running on the real checkpoints: 3,988,052 PBR voxels x 6 channels in 65.8s (base_color 0:3, metallic 3:4, roughness 4:5, alpha 5:6). Two things here fail SILENTLY rather than loudly, so both are asserted in comments: 1. shape_slat arrives DENORMALISED - the shape stage un-standardises it for the decoder - but the texture flow was trained against the standardised form. It is re-normalised before use as concat_cond. Skipping that gives a plausible mesh with wrong colours, not an error. 2. tex_dec has pred_subdiv=False: it cannot invent subdivisions and must be handed the shape decoder's subs as guides, so texture voxels land on the geometry that was actually built. The decoder's output is mapped * 0.5 + 0.5 into [0,1], the range o_voxel expects. BAKE ORDER. Handing o_voxel the raw ~8M-face mesh hangs - the same wall the standalone remesh test hit (killed at 20min), and the trellis2 lane's own operator note says the uncapped bake peaks at 75GB. So the mesh is welded, stripped of floaters and decimated BEFORE baking; the baker samples the attribute VOLUME at mesh positions, so a decimated mesh still gets correct colours. Measured on the way through: welded 3,988,052 -> 3,983,672 verts floaters 12 components -> 1 kept, 6,332 faces dropped decimated 7,996,876 -> 214,322 faces pre-bake 34.6s That floater count is worth noting: 12 components, not the 52,855 the first health pass reported. Welding first is what makes the difference. remesh now defaults OFF in to_glb, unlike upstream. Upstream runs on CUDA; this is the CPU/Metal build and its remesher took >20 minutes on a 214k-face mesh. It is also handed an already-clean mesh, so there is far less for it to fix. Operator gains texture + texture_size params; geometry-only stays the default because it is ~3min against the textured path's extra flow and bake. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
78 lines
3.0 KiB
Python
78 lines
3.0 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",
|
|
"slat_tex": "slat_flow_imgshape2tex_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"
|
|
|
|
flows = [("ss_flow", ss_flow.load), ("ss_dec", ss_dec.load),
|
|
("slat_512", slat_flow.load), ("slat_1024", slat_flow.load)]
|
|
if with_texture:
|
|
flows.append(("slat_tex", slat_flow.load))
|
|
|
|
models = {}
|
|
for key, loader in flows:
|
|
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)
|
|
if with_texture:
|
|
# tex_1024 differs from shape_1024 only in naf_target_size (1024 vs 512), but
|
|
# that changes the high-res branch it samples, so it needs its own conditioner
|
|
models["cond_tex"] = ProjConditioner("tex_1024", device=device)
|
|
return models
|