18179bc479
All seven Pixal3D models now load and run. shape_dec 292/292, tex_dec 284/284, both with zero missing/unmapped/mismatched keys - the 8-param difference between them is exactly the four to_subdiv layers, since tex_dec has pred_subdiv=False. Behaviour is right: 12 latent voxels grow SELECTIVELY through four stages (12 -> 91 -> 193 -> 358 -> 1150) rather than x8 each time, which would have reached 49,152. Scale lands at 1/16. shape_dec emits the hardcoded 7 channels and its vertex head produces offsets inside the [-0.5, 1.5] band its sigmoid+voxel_margin allows. tex_dec, guided by shape_dec's masks, reproduces exactly the same voxel count. VERIFICATION CAVEAT, recorded in the README: these two are the only models using sparse conv, so upstream cannot run here and there is no numerical oracle. Unlike the five models verified at correlation 1.0, these are checked structurally and behaviourally only. Weaker evidence, and labelled as such rather than presented alongside the verified results.
111 lines
4.0 KiB
Python
111 lines
4.0 KiB
Python
"""Behavioural tests for the two sparse decoders.
|
|
|
|
These are the only Pixal3D models using sparse convolution, so upstream cannot run here
|
|
(spconv has no Metal build) and there is NO numerical oracle — unlike the flow models and
|
|
ss_dec, which are all verified at correlation 1.0 against upstream on CPU torch.
|
|
|
|
So these check structure and behaviour rather than values: that weights map completely,
|
|
that the occupied set grows SELECTIVELY rather than x8, that scale tracks the four
|
|
upsamples, and that the vertex head lands inside the band its sigmoid+margin allows.
|
|
Weaker than a correlation, and stated as such.
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
from pixal3d_mlx.decoders import load # noqa: E402
|
|
from trellis_sparse_mlx import SparseTensor # noqa: E402
|
|
|
|
CK = Path("/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts")
|
|
W = Path(__file__).resolve().parents[1] / "weights"
|
|
|
|
|
|
def _latent(n=12, ch=32, res=4, seed=0):
|
|
rng = np.random.default_rng(seed)
|
|
co = np.concatenate(
|
|
[np.zeros((n, 1), dtype=np.int32), rng.integers(0, res, (n, 3)).astype(np.int32)], 1
|
|
)
|
|
co = np.unique(co, axis=0)
|
|
return SparseTensor(
|
|
mx.array(rng.standard_normal((len(co), ch)).astype(np.float32)), mx.array(co)
|
|
)
|
|
|
|
|
|
def _have(stem):
|
|
return (W / f"{stem}.safetensors").exists() and (CK / f"{stem}.json").exists()
|
|
|
|
|
|
def test_shape_dec():
|
|
stem = "shape_dec_next_dc_f16c32_fp16"
|
|
if not _have(stem):
|
|
print(" (skipped: weights absent)")
|
|
return 0.0
|
|
m, rep = load(W / f"{stem}.safetensors", CK / f"{stem}.json")
|
|
assert not rep["missing"] and not rep["unmapped"] and not rep["mismatched"], rep
|
|
x = _latent()
|
|
n_in = x.coords.shape[0]
|
|
out, subs = m(x, return_subs=True)
|
|
f = np.asarray(out.feats)
|
|
|
|
assert f.shape[1] == 7, f"FlexiDualGrid emits 7 channels, got {f.shape[1]}"
|
|
assert np.isfinite(f).all(), "non-finite output"
|
|
# four x2 upsamples -> scale 1/16
|
|
assert abs(out._scale[0] - 1 / 16) < 1e-9, out._scale
|
|
# growth must be selective, not the full x8 per stage
|
|
assert out.coords.shape[0] > n_in, "no growth at all"
|
|
assert out.coords.shape[0] < n_in * 8**4, "grew by the full x8 every stage"
|
|
assert len(subs) == 4, f"expected 4 subdivision masks, got {len(subs)}"
|
|
|
|
v = np.asarray(m.decode_vertices(out))
|
|
m_ = m.voxel_margin
|
|
assert v.min() >= -m_ - 1e-4 and v.max() <= 1 + m_ + 1e-4, (v.min(), v.max())
|
|
return 0.0
|
|
|
|
|
|
def test_tex_dec_follows_guide_masks():
|
|
"""tex_dec has pred_subdiv=False, so its structure must match the masks it is given."""
|
|
s_stem, t_stem = "shape_dec_next_dc_f16c32_fp16", "tex_dec_next_dc_f16c32_fp16"
|
|
if not (_have(s_stem) and _have(t_stem)):
|
|
print(" (skipped: weights absent)")
|
|
return 0.0
|
|
sm, _ = load(W / f"{s_stem}.safetensors", CK / f"{s_stem}.json")
|
|
tm, rep = load(W / f"{t_stem}.safetensors", CK / f"{t_stem}.json")
|
|
assert not rep["missing"] and not rep["unmapped"], rep
|
|
|
|
x = _latent()
|
|
shape_out, subs = sm(x, return_subs=True)
|
|
tex_out = tm(_latent(), guide_subs=subs)
|
|
f = np.asarray(tex_out.feats)
|
|
|
|
assert f.shape[1] == 6, f"tex_dec emits 6 channels, got {f.shape[1]}"
|
|
assert np.isfinite(f).all(), "non-finite output"
|
|
assert tex_out.coords.shape[0] == shape_out.coords.shape[0], (
|
|
"guided decoder produced a different voxel count than the masks describe"
|
|
)
|
|
return 0.0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tests = [
|
|
("shape_dec structure", test_shape_dec),
|
|
("tex_dec follows guides", test_tex_dec_follows_guide_masks),
|
|
]
|
|
failed = 0
|
|
for name, fn in tests:
|
|
try:
|
|
fn()
|
|
print(f" PASS {name}")
|
|
except AssertionError as e:
|
|
print(f" FAIL {name}: {e}")
|
|
failed += 1
|
|
except Exception as e: # noqa: BLE001
|
|
print(f" ERROR {name}: {type(e).__name__}: {e}")
|
|
failed += 1
|
|
print(f"\n{len(tests)-failed}/{len(tests)} passed")
|
|
sys.exit(1 if failed else 0)
|