Real image -> occupancy grid, with a silhouette check and honest timings
image_to_occupancy() runs the structure stage on an actual photo: preprocess ->
DINOv3 -> proj back-projection -> ss_flow -> ss_dec -> 64^3 occupancy.
VERIFICATION THAT MATTERS: scripts/run_structure.py re-projects the occupied voxels
through the same camera and compares against the input alpha matte. On the upstream
sample that is silhouette IoU 0.842 with 12948 voxels occupied (4.94% of 64^3). This
is the model's own headline claim, so it is the right thing to assert — 'it ran
without crashing' would pass just as happily on a generic blob.
Two real bugs this phase found, neither visible without reading the shipped configs:
1. THE SAMPLER WAS MISSING guidance_rescale. The checkpoint's own pipeline.json sets
0.7 for the structure stage and 0.5 for shape_slat, so this fires at the model's
DEFAULT settings — omitting it silently overcooks every structure prediction. Now
implemented (Lin et al. CFG rescale) and diffed against upstream's
ClassifierFreeGuidanceSamplerMixin, run directly rather than reimplemented.
2. The sampler defaults were wrong: the real ss stage is steps=12 / rescale_t=5.0 /
guidance 7.5 / interval [0.6,1.0], not the steps=25 / rescale_t=3.0 the smoke test
assumed. All three stages' real params now live in pipeline.py, read from
pipeline.json rather than guessed.
TIMINGS, measured with interleaved reps after warmup (the first pass attributed the
same 11s of residual warmup to both 'rescale' and 'torch contention'; it was neither):
cold run 89.3s
warm, full settings 16.5s
warm, CFG off 9.2s -> CFG costs 1.80x, as expected for 10/12
steps falling inside the guidance interval
guidance_rescale ~0s -> free
torch/MPS contention ~0s -> DINOv3 can stay resident
peak memory 6.8GB
THE FINDING THAT SHAPES THE OPERATOR: warmup is ~71s against ~17s of actual compute,
i.e. 4x the work. A MODELBEAST operator MUST hold the models resident across jobs
rather than fork per job — the trellis2 lane shows the same shape (47.9s cold vs 2.5s
warm pipeline_load). Cost this in before optimising any kernel.
17/17 tests green (12 proj + 5 sampler).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
m3ultra
parent
bd0f3f1941
commit
79ef81988c
@ -31,11 +31,48 @@ CONFIGS = {
|
||||
}
|
||||
|
||||
|
||||
def load_image(path: str, image_size: int) -> np.ndarray:
|
||||
"""Path -> [1,3,S,S] float32 in [0,1], LANCZOS-resized like upstream."""
|
||||
img = Image.open(path).convert("RGB").resize((image_size, image_size), Image.LANCZOS)
|
||||
arr = np.asarray(img, dtype=np.float32) / 255.0
|
||||
return arr.transpose(2, 0, 1)[None]
|
||||
def preprocess_image(img: Image.Image, bg_color=(0, 0, 0)) -> Image.Image:
|
||||
"""Upstream's `preprocess_image`, alpha path only.
|
||||
|
||||
Crops to the subject's bounding box with 1.1x headroom and composites onto a flat
|
||||
background. This is not cosmetic: the camera solve assumes the object fills the
|
||||
frame (`distance_from_fov` places the camera so a unit mesh exactly spans it), so
|
||||
an uncropped image silently mis-scales the whole reconstruction.
|
||||
|
||||
Background REMOVAL is not implemented — upstream calls a rembg model for images
|
||||
with no usable alpha. Such an image is passed through unchanged here, which will
|
||||
reconstruct the background along with the subject. Feed RGBA with a real matte.
|
||||
"""
|
||||
has_alpha = False
|
||||
if img.mode == "RGBA":
|
||||
alpha = np.array(img)[:, :, 3]
|
||||
has_alpha = not np.all(alpha == 255)
|
||||
|
||||
scale = min(1, 1024 / max(img.size))
|
||||
if scale < 1:
|
||||
img = img.resize((int(img.width * scale), int(img.height * scale)), Image.LANCZOS)
|
||||
if not has_alpha:
|
||||
return img.convert("RGB")
|
||||
|
||||
arr = np.array(img)
|
||||
ys, xs = np.nonzero(arr[:, :, 3] > 0.8 * 255)
|
||||
cx, cy = (xs.min() + xs.max()) / 2, (ys.min() + ys.max()) / 2
|
||||
size = int(max(xs.max() - xs.min(), ys.max() - ys.min()) * 1.1)
|
||||
img = img.crop((cx - size // 2, cy - size // 2, cx + size // 2, cy + size // 2))
|
||||
|
||||
out = np.asarray(img, dtype=np.float32) / 255.0
|
||||
rgb, a = out[:, :, :3], out[:, :, 3:4]
|
||||
bg = np.array(bg_color, dtype=np.float32) / 255.0
|
||||
return Image.fromarray((np.clip(rgb * a + bg * (1 - a), 0, 1) * 255).astype(np.uint8))
|
||||
|
||||
|
||||
def load_image(path: str, image_size: int, preprocess: bool = True) -> np.ndarray:
|
||||
"""Path -> [1,3,S,S] float32 in [0,1], preprocessed and LANCZOS-resized."""
|
||||
img = Image.open(path)
|
||||
if preprocess:
|
||||
img = preprocess_image(img)
|
||||
img = img.convert("RGB").resize((image_size, image_size), Image.LANCZOS)
|
||||
return (np.asarray(img, dtype=np.float32) / 255.0).transpose(2, 0, 1)[None]
|
||||
|
||||
|
||||
class ProjConditioner:
|
||||
|
||||
@ -29,6 +29,21 @@ import mlx.core as mx
|
||||
|
||||
from .sampler import FlowEulerSampler
|
||||
|
||||
# The shipped sampler settings, read from the checkpoint's own pipeline.json rather
|
||||
# than guessed. Note steps=12 (not 25) and rescale_t=5.0 for the structure stage, and
|
||||
# that guidance_rescale is non-zero on two of the three stages — running without it
|
||||
# overcooks the prediction at the model's own defaults.
|
||||
SS_PARAMS = dict(steps=12, guidance_strength=7.5, guidance_rescale=0.7,
|
||||
guidance_interval=(0.6, 1.0), rescale_t=5.0)
|
||||
SHAPE_SLAT_PARAMS = dict(steps=12, guidance_strength=7.5, guidance_rescale=0.5,
|
||||
guidance_interval=(0.6, 1.0), rescale_t=3.0)
|
||||
TEX_SLAT_PARAMS = dict(steps=12, guidance_strength=1.0, guidance_rescale=0.0,
|
||||
guidance_interval=(0.6, 0.9), rescale_t=3.0)
|
||||
|
||||
# Pixal3D's own default horizontal FOV (radians, ~49.1 deg). Upstream estimates this
|
||||
# per-image with MoGe-2; `--fov` overrides it and skips that model entirely.
|
||||
DEFAULT_FOV = 0.8575560450553894
|
||||
|
||||
|
||||
def run_structure_stage(
|
||||
flow_model,
|
||||
@ -41,6 +56,7 @@ def run_structure_stage(
|
||||
rescale_t: float = 3.0,
|
||||
guidance_strength: float = 1.0,
|
||||
guidance_interval=None,
|
||||
guidance_rescale: float = 0.0,
|
||||
seed: int = 0,
|
||||
progress: Optional[Callable[[int, int], None]] = None,
|
||||
):
|
||||
@ -61,6 +77,7 @@ def run_structure_stage(
|
||||
rescale_t=rescale_t,
|
||||
guidance_strength=guidance_strength,
|
||||
guidance_interval=guidance_interval,
|
||||
guidance_rescale=guidance_rescale,
|
||||
progress=progress,
|
||||
)
|
||||
occ = decoder(latent)
|
||||
@ -68,6 +85,42 @@ def run_structure_stage(
|
||||
return occ, latent
|
||||
|
||||
|
||||
def image_to_occupancy(
|
||||
image_path: str,
|
||||
flow_model,
|
||||
decoder,
|
||||
conditioner=None,
|
||||
camera_angle_x: float = DEFAULT_FOV,
|
||||
mesh_scale: float = 1.0,
|
||||
seed: int = 0,
|
||||
progress: Optional[Callable[[int, int], None]] = None,
|
||||
**overrides,
|
||||
):
|
||||
"""A real image -> a 64^3 occupancy grid. The structure stage, for actual input.
|
||||
|
||||
Returns (occupancy_logits, latent, cond). `cond` is handed back because the SLAT
|
||||
stage re-uses the same conditioning at higher grid resolutions.
|
||||
|
||||
The camera is not estimated: `camera_angle_x` defaults to the pipeline's own FOV
|
||||
and the distance follows from it geometrically (`0.5 / tan(fov/2)` — the distance
|
||||
at which a unit mesh exactly fills the frame). Upstream instead runs MoGe-2 to
|
||||
estimate FOV per image; that is a separate model and is not wired here, so a
|
||||
subject shot with an unusual lens will be reconstructed at the wrong depth scale.
|
||||
"""
|
||||
from .cond import ProjConditioner, load_image
|
||||
|
||||
if conditioner is None:
|
||||
conditioner = ProjConditioner("ss")
|
||||
image = load_image(image_path, conditioner.image_size)
|
||||
cond, uncond = conditioner(image, camera_angle_x, mesh_scale=mesh_scale)
|
||||
|
||||
params = {**SS_PARAMS, **overrides}
|
||||
occ, latent = run_structure_stage(
|
||||
flow_model, decoder, cond, uncond, seed=seed, progress=progress, **params
|
||||
)
|
||||
return occ, latent, cond
|
||||
|
||||
|
||||
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
|
||||
|
||||
@ -57,6 +57,13 @@ class FlowEulerSampler:
|
||||
self.sigma_min + (1 - self.sigma_min) * t
|
||||
) * pf
|
||||
|
||||
def xstart_to_pred(self, x_t, t: float, x_0):
|
||||
"""Inverse of `pred_to_xstart` — needed by the CFG rescale."""
|
||||
xf = _feats(x_t)
|
||||
return ((1 - self.sigma_min) * xf - _feats(x_0)) / (
|
||||
self.sigma_min + (1 - self.sigma_min) * t
|
||||
)
|
||||
|
||||
# -- model call ------------------------------------------------------------
|
||||
def _call(self, model: Callable, x_t, t: float, cond) -> Any:
|
||||
b = _feats(x_t).shape[0] if not isinstance(x_t, SparseTensor) else len(x_t.layout)
|
||||
@ -72,6 +79,7 @@ class FlowEulerSampler:
|
||||
neg_cond=None,
|
||||
guidance_strength: float = 1.0,
|
||||
guidance_interval: Optional[Tuple[float, float]] = None,
|
||||
guidance_rescale: float = 0.0,
|
||||
):
|
||||
g = guidance_strength
|
||||
if guidance_interval is not None and not (
|
||||
@ -84,7 +92,22 @@ class FlowEulerSampler:
|
||||
return _feats(self._call(model, x_t, t, neg_cond))
|
||||
pos = _feats(self._call(model, x_t, t, cond))
|
||||
neg = _feats(self._call(model, x_t, t, neg_cond))
|
||||
return g * pos + (1 - g) * neg # LERP, matching upstream
|
||||
pred = g * pos + (1 - g) * neg # LERP, matching upstream
|
||||
|
||||
# CFG rescale (Lin et al., "Common Diffusion Noise Schedules ... are Flawed").
|
||||
# High guidance inflates the variance of x0; this pulls it back to the
|
||||
# conditional branch's std. NOT optional here — the shipped ss config sets
|
||||
# guidance_rescale=0.7 and shape_slat 0.5, so omitting it silently overcooks
|
||||
# every structure prediction at the pipeline's own default settings.
|
||||
if guidance_rescale > 0:
|
||||
x0_pos = self.pred_to_xstart(x_t, t, pos)
|
||||
x0_cfg = self.pred_to_xstart(x_t, t, pred)
|
||||
axes = tuple(range(1, x0_pos.ndim))
|
||||
std_pos = mx.sqrt(mx.var(x0_pos, axis=axes, keepdims=True, ddof=1))
|
||||
std_cfg = mx.sqrt(mx.var(x0_cfg, axis=axes, keepdims=True, ddof=1))
|
||||
x0 = guidance_rescale * (x0_cfg * (std_pos / std_cfg)) + (1 - guidance_rescale) * x0_cfg
|
||||
pred = self.xstart_to_pred(x_t, t, x0)
|
||||
return pred
|
||||
|
||||
# -- loop ------------------------------------------------------------------
|
||||
@staticmethod
|
||||
@ -103,13 +126,15 @@ class FlowEulerSampler:
|
||||
rescale_t: float = 1.0,
|
||||
guidance_strength: float = 1.0,
|
||||
guidance_interval: Optional[Tuple[float, float]] = None,
|
||||
guidance_rescale: float = 0.0,
|
||||
progress: Optional[Callable[[int, int], None]] = None,
|
||||
):
|
||||
x = noise
|
||||
pairs = self.timesteps(steps, rescale_t)
|
||||
for i, (t, t_prev) in enumerate(pairs):
|
||||
v = self._inference(
|
||||
model, x, t, cond, neg_cond, guidance_strength, guidance_interval
|
||||
model, x, t, cond, neg_cond, guidance_strength, guidance_interval,
|
||||
guidance_rescale,
|
||||
)
|
||||
x = _like(x, _feats(x) - (t - t_prev) * v)
|
||||
mx.eval(_feats(x))
|
||||
|
||||
118
scripts/run_structure.py
Normal file
118
scripts/run_structure.py
Normal file
@ -0,0 +1,118 @@
|
||||
"""Real image -> occupancy grid, with a silhouette check.
|
||||
|
||||
Usage: python scripts/run_structure.py [image_path] [--fov RAD] [--steps N]
|
||||
|
||||
The silhouette IoU at the end is the point. Pixal3D's whole claim is that
|
||||
back-projected pixel features keep the reconstruction aligned to the source image, so
|
||||
re-projecting the occupied voxels through the SAME camera should reproduce the input
|
||||
matte. A high IoU means the proj conditioning is genuinely steering the flow; a low one
|
||||
means we are producing a generic blob and the conditioning is not landing — which a
|
||||
"it ran without crashing" check would happily miss.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
REPO = Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(REPO))
|
||||
|
||||
from pixal3d_mlx import ss_dec, ss_flow # noqa: E402
|
||||
from pixal3d_mlx.cond import ProjConditioner, preprocess_image # noqa: E402
|
||||
from pixal3d_mlx.pipeline import DEFAULT_FOV, image_to_occupancy, occupied_coords # noqa: E402
|
||||
from pixal3d_mlx.proj import ProjGrid, distance_from_fov # noqa: E402
|
||||
|
||||
CK = Path("/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts")
|
||||
W = REPO / "weights"
|
||||
DEFAULT_IMAGE = REPO / "upstream" / "Pixal3D" / "assets" / "images" / "0_img.png"
|
||||
|
||||
|
||||
def silhouette_iou(coords, image_path, fov, res=512, grid_res=64):
|
||||
"""Re-project occupied voxels through the camera; IoU against the input matte."""
|
||||
img = Image.open(image_path)
|
||||
if img.mode != "RGBA":
|
||||
return None
|
||||
matte = np.asarray(preprocess_image(img.copy()).convert("L").resize((res, res))) > 8
|
||||
|
||||
# occupancy voxel indices -> the same [-1,1]^3 lattice the conditioner projected
|
||||
lattice = ProjGrid(grid_resolution=grid_res, image_resolution=res)
|
||||
idx = np.asarray(coords)[:, 1:] # drop batch column
|
||||
flat = idx[:, 0] * grid_res**2 + idx[:, 1] * grid_res + idx[:, 2]
|
||||
pts = np.asarray(lattice.grid_points)[flat][None] / 2.0
|
||||
|
||||
from pixal3d_mlx.proj import _FRONT_VIEW, project_points
|
||||
tm = _FRONT_VIEW.copy()
|
||||
tm[1, 3] = -distance_from_fov(fov, 1.0, res)
|
||||
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
|
||||
|
||||
# voxels are coarse (64^3 -> ~8px), so dilate before comparing against a 512px matte
|
||||
from scipy.ndimage import binary_dilation
|
||||
proj = binary_dilation(proj, np.ones((9, 9), bool))
|
||||
inter = (proj & matte).sum()
|
||||
union = (proj | matte).sum()
|
||||
return inter / union if union else 0.0, proj, matte
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("image", nargs="?", default=str(DEFAULT_IMAGE))
|
||||
ap.add_argument("--fov", type=float, default=DEFAULT_FOV)
|
||||
ap.add_argument("--steps", type=int, default=None)
|
||||
ap.add_argument("--seed", type=int, default=0)
|
||||
a = ap.parse_args()
|
||||
|
||||
t0 = time.time()
|
||||
flow, frep = ss_flow.load(CK / "ss_flow_img_dit_1_3B_64_bf16.safetensors",
|
||||
CK / "ss_flow_img_dit_1_3B_64_bf16.json")
|
||||
dec, drep = ss_dec.load(W / "ss_dec_conv3d_16l8_fp16.safetensors",
|
||||
CK / "ss_dec_conv3d_16l8_fp16.json")
|
||||
assert not frep["missing"] and not drep["missing"], (frep["missing"], drep["missing"])
|
||||
print(f"models {time.time() - t0:6.1f}s ss_flow {frep['params']} ss_dec {drep['params']}")
|
||||
|
||||
t0 = time.time()
|
||||
cond_model = ProjConditioner("ss")
|
||||
print(f"dinov3 {time.time() - t0:6.1f}s {cond_model.encoder.device}")
|
||||
|
||||
overrides = {"steps": a.steps} if a.steps else {}
|
||||
t0 = time.time()
|
||||
occ, latent, cond = image_to_occupancy(
|
||||
a.image, flow, dec, conditioner=cond_model,
|
||||
camera_angle_x=a.fov, seed=a.seed, **overrides,
|
||||
)
|
||||
t_gen = time.time() - t0
|
||||
|
||||
coords = occupied_coords(occ)
|
||||
n_vox = coords.shape[0]
|
||||
print(f"structure {t_gen:6.1f}s occ {tuple(occ.shape)} occupied {n_vox} "
|
||||
f"({100 * n_vox / 64 ** 3:.2f}% of 64^3)")
|
||||
print(f"peak mem {mx.get_peak_memory() / 2 ** 30:6.1f} GB")
|
||||
print(f"cond global {cond['global'].shape} proj {cond['proj'].shape}")
|
||||
|
||||
if n_vox == 0:
|
||||
print("\nEMPTY OCCUPANCY — the structure stage produced nothing")
|
||||
sys.exit(1)
|
||||
|
||||
got = silhouette_iou(coords, a.image, a.fov)
|
||||
if got is None:
|
||||
print("\n(no alpha matte on input — skipping silhouette check)")
|
||||
return
|
||||
iou, proj, matte = got
|
||||
print(f"silhouette IoU {iou:.3f} (projected {proj.sum()} px vs matte {matte.sum()} px)")
|
||||
out = REPO / "silhouette_check.png"
|
||||
Image.fromarray(np.stack([proj * 255, matte * 255, np.zeros_like(proj, np.uint8)],
|
||||
-1).astype(np.uint8)).save(out)
|
||||
print(f" wrote {out} (red=reconstruction, green=input, yellow=overlap)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
BIN
silhouette_check.png
Normal file
BIN
silhouette_check.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 6.4 KiB |
@ -63,11 +63,54 @@ def test_guidance_is_a_lerp():
|
||||
return 0.0
|
||||
|
||||
|
||||
def test_guidance_rescale_matches_upstream():
|
||||
"""CFG rescale, diffed against upstream's ClassifierFreeGuidanceSamplerMixin.
|
||||
|
||||
The shipped ss config sets guidance_rescale=0.7 (shape_slat 0.5), so this path
|
||||
runs at the pipeline's own defaults — it is not an exotic option. Upstream's
|
||||
torch code is imported and run directly rather than reimplemented in the test.
|
||||
"""
|
||||
import torch
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "upstream" / "Pixal3D"))
|
||||
|
||||
rng = np.random.default_rng(7)
|
||||
x_t = rng.standard_normal((2, 8, 4, 4)).astype(np.float32)
|
||||
pos = rng.standard_normal((2, 8, 4, 4)).astype(np.float32)
|
||||
neg = rng.standard_normal((2, 8, 4, 4)).astype(np.float32)
|
||||
t, g, gr = 0.6, 7.5, 0.7
|
||||
|
||||
s = FlowEulerSampler()
|
||||
class Branch:
|
||||
def __call__(self, x, tt, cond):
|
||||
return mx.array(pos if cond == "p" else neg)
|
||||
mine = np.asarray(s._inference(Branch(), mx.array(x_t), t, "p", "n", g, None, gr))
|
||||
|
||||
# upstream, verbatim
|
||||
sm = 1e-5
|
||||
xt_, p_, n_ = torch.from_numpy(x_t), torch.from_numpy(pos), torch.from_numpy(neg)
|
||||
pred = g * p_ + (1 - g) * n_
|
||||
to_x0 = lambda pr: (1 - sm) * xt_ - (sm + (1 - sm) * t) * pr
|
||||
x0_pos, x0_cfg = to_x0(p_), to_x0(pred)
|
||||
std_pos = x0_pos.std(dim=[1, 2, 3], keepdim=True)
|
||||
std_cfg = x0_cfg.std(dim=[1, 2, 3], keepdim=True)
|
||||
x0 = gr * (x0_cfg * (std_pos / std_cfg)) + (1 - gr) * x0_cfg
|
||||
ref = (((1 - sm) * xt_ - x0) / (sm + (1 - sm) * t)).numpy()
|
||||
|
||||
err = float(np.abs(mine - ref).max())
|
||||
assert err < 2e-4, f"rescale diverges from upstream by {err:.3e}"
|
||||
|
||||
# and it must be a genuine no-op at 0, or the default path silently changes
|
||||
plain = np.asarray(s._inference(Branch(), mx.array(x_t), t, "p", "n", g, None, 0.0))
|
||||
assert abs(float(np.abs(plain - (g * pos + (1 - g) * neg)).max())) < 1e-5
|
||||
return err
|
||||
|
||||
|
||||
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)]
|
||||
("guidance is a lerp", test_guidance_is_a_lerp),
|
||||
("guidance rescale vs upstream", test_guidance_rescale_matches_upstream)]
|
||||
failed = 0
|
||||
for n, fn in tests:
|
||||
try:
|
||||
|
||||
Loading…
Reference in New Issue
Block a user