diff --git a/pixal3d_mlx/cond.py b/pixal3d_mlx/cond.py index ea54b4f..92f6ff7 100644 --- a/pixal3d_mlx/cond.py +++ b/pixal3d_mlx/cond.py @@ -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: diff --git a/pixal3d_mlx/pipeline.py b/pixal3d_mlx/pipeline.py index 6342b14..b9cc9d9 100644 --- a/pixal3d_mlx/pipeline.py +++ b/pixal3d_mlx/pipeline.py @@ -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 diff --git a/pixal3d_mlx/sampler.py b/pixal3d_mlx/sampler.py index 3f4116c..a9f9fd5 100644 --- a/pixal3d_mlx/sampler.py +++ b/pixal3d_mlx/sampler.py @@ -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)) diff --git a/scripts/run_structure.py b/scripts/run_structure.py new file mode 100644 index 0000000..aec4488 --- /dev/null +++ b/scripts/run_structure.py @@ -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() diff --git a/silhouette_check.png b/silhouette_check.png new file mode 100644 index 0000000..7a89089 Binary files /dev/null and b/silhouette_check.png differ diff --git a/tests/test_sampler.py b/tests/test_sampler.py index 55e605d..d817406 100644 --- a/tests/test_sampler.py +++ b/tests/test_sampler.py @@ -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: