d43a6dfd69
The structure stage now runs end to end on Metal: noise -> ss_flow (12-step Euler) -> latent -> ss_dec -> 64^3 occupancy grid. 759 ms/step for the 1.3B DiT, 50 ms for the decoder. Output lands at 1.36% occupancy, which is the right order for a surface in a 64^3 grid. Sampler details worth recording, both from upstream: - CFG is a LERP (g*pos + (1-g)*neg), NOT neg + g*(pos-neg). Those differ non-linearly in strength rather than failing outright, so it is silent when wrong. Tested. - Guidance interval forces strength to 1 outside its window, which halves model calls there: 6 calls for 4 steps rather than 8. Tested by counting. Schedule matches upstream to 1e-16 and constant velocity integrates exactly at any step count. Remaining for a real image->3D run is CONDITIONING, not models: DINOv3 features plus Pixal3D's camera back-projection for the view-aligned 'proj' half. Deliberately not porting DINOv3 - it is a stock ViT run once per image, outside the denoising loop, so torch on MPS is the right tool and transformers gives exact parity for free.
79 lines
3.0 KiB
Python
79 lines
3.0 KiB
Python
"""Flow Euler sampler tests — schedule, integration, and guidance call counts."""
|
|
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.sampler import FlowEulerSampler # noqa: E402
|
|
|
|
|
|
def test_schedule_matches_upstream():
|
|
s = FlowEulerSampler()
|
|
worst = 0.0
|
|
for r in (1.0, 3.0, 0.5):
|
|
got = s.timesteps(7, r)
|
|
seq = np.linspace(1, 0, 8)
|
|
seq = r * seq / (1 + (r - 1) * seq)
|
|
want = [(seq[i], seq[i + 1]) for i in range(7)]
|
|
worst = max(worst, max(abs(a - c) + abs(b - d) for (a, b), (c, d) in zip(got, want)))
|
|
assert worst < 1e-12, f"schedule err {worst:.3g}"
|
|
return worst
|
|
|
|
|
|
def test_constant_velocity_integrates_exactly():
|
|
"""v constant => x_end = x_start - v*(1-0), regardless of step count."""
|
|
s = FlowEulerSampler()
|
|
class Const:
|
|
def __call__(self, x, t, cond):
|
|
return mx.ones_like(x) * 2.0
|
|
for steps in (1, 10, 37):
|
|
out = np.asarray(s.sample(Const(), mx.zeros((1, 4)), cond=None, steps=steps))
|
|
err = abs(out[0, 0] - (-2.0))
|
|
assert err < 1e-5, f"steps={steps} gave {out[0,0]}, want -2.0"
|
|
return 0.0
|
|
|
|
|
|
def test_guidance_interval_skips_negative_pass():
|
|
"""Outside the interval only the conditional branch runs — half the model calls."""
|
|
s = FlowEulerSampler()
|
|
calls = []
|
|
class Counting:
|
|
def __call__(self, x, t, cond):
|
|
calls.append(float(t[0]) / 1000.0)
|
|
return mx.zeros_like(x)
|
|
s.sample(Counting(), mx.zeros((1, 4)), cond="p", neg_cond="n", steps=4,
|
|
guidance_strength=3.0, guidance_interval=(0.0, 0.5))
|
|
assert len(calls) == 6, f"expected 6 calls (2 in-window + 1 out), got {len(calls)}"
|
|
from collections import Counter
|
|
c = Counter(round(t, 3) for t in calls)
|
|
assert c[1.0] == 1 and c[0.5] == 2, dict(c)
|
|
return 0.0
|
|
|
|
|
|
def test_guidance_is_a_lerp():
|
|
"""upstream uses g*pos + (1-g)*neg, not neg + g*(pos-neg)."""
|
|
s = FlowEulerSampler()
|
|
class Two:
|
|
def __call__(self, x, t, cond):
|
|
return mx.ones_like(x) * (1.0 if cond == "p" else 3.0)
|
|
out = np.asarray(s.sample(Two(), mx.zeros((1, 1)), cond="p", neg_cond="n",
|
|
steps=1, guidance_strength=2.0))
|
|
# g*1 + (1-g)*3 = 2 - 3 = -1 -> x = 0 - 1*(-1) = +1
|
|
assert abs(out[0, 0] - 1.0) < 1e-5, f"got {out[0,0]}, want +1.0 for a LERP"
|
|
return 0.0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tests = [("schedule vs upstream", test_schedule_matches_upstream),
|
|
("constant v integrates", test_constant_velocity_integrates_exactly),
|
|
("guidance interval", test_guidance_interval_skips_negative_pass),
|
|
("guidance is a lerp", test_guidance_is_a_lerp)]
|
|
failed = 0
|
|
for n, fn in tests:
|
|
try:
|
|
fn(); print(f" PASS {n}")
|
|
except AssertionError as e:
|
|
print(f" FAIL {n}: {e}"); failed += 1
|
|
print(f"\n{len(tests)-failed}/{len(tests)} passed")
|
|
sys.exit(1 if failed else 0)
|