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.
78 lines
2.6 KiB
Python
78 lines
2.6 KiB
Python
"""Pixal3D stage wiring for MLX.
|
|
|
|
What runs today: the STRUCTURE stage, end to end on Metal —
|
|
noise -> ss_flow (Euler sampling) -> latent -> ss_dec -> occupancy grid
|
|
|
|
What is still missing for a real image->3D run is the conditioning, not the models.
|
|
Pixal3D conditions on DINOv3 features in two forms:
|
|
|
|
global : [B, M, 1024] pooled image tokens, cross-attended
|
|
proj : [B, N, 1024] VIEW-ALIGNED features, one token per spatial position,
|
|
projected and ADDED rather than attended
|
|
|
|
The `proj` half is the actual novelty — it back-projects pixel features into 3D through
|
|
the camera, which is what keeps silhouettes exact. Producing it needs the DINOv3 encoder
|
|
plus Pixal3D's `DinoV3ProjFeatureExtractor` camera logic.
|
|
|
|
Deliberate design note: DINOv3 is a stock ViT-L/16 run ONCE per image, not inside the
|
|
25-step denoising loop, so it is not worth porting to MLX — torch on MPS runs it fine and
|
|
using `transformers` gives exact parity with upstream for free. The effort belongs in the
|
|
30-block DiT loops, which are already here and verified. Note facebook/dinov3 is gated;
|
|
upstream itself falls back to the ungated `camenduru` mirror.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Callable, Optional
|
|
|
|
import mlx.core as mx
|
|
|
|
from .sampler import FlowEulerSampler
|
|
|
|
|
|
def run_structure_stage(
|
|
flow_model,
|
|
decoder,
|
|
cond,
|
|
neg_cond=None,
|
|
resolution: int = 16,
|
|
latent_channels: int = 8,
|
|
steps: int = 25,
|
|
rescale_t: float = 3.0,
|
|
guidance_strength: float = 1.0,
|
|
guidance_interval=None,
|
|
seed: int = 0,
|
|
progress: Optional[Callable[[int, int], None]] = None,
|
|
):
|
|
"""noise -> ss_flow -> ss_dec. Returns (occupancy_logits, latent).
|
|
|
|
`occupancy_logits` is [B, 1, R*4, R*4, R*4]; threshold at 0 for the occupied set that
|
|
seeds the SLAT stage.
|
|
"""
|
|
mx.random.seed(seed)
|
|
noise = mx.random.normal((1, latent_channels, resolution, resolution, resolution))
|
|
sampler = FlowEulerSampler()
|
|
latent = sampler.sample(
|
|
flow_model,
|
|
noise,
|
|
cond=cond,
|
|
neg_cond=neg_cond,
|
|
steps=steps,
|
|
rescale_t=rescale_t,
|
|
guidance_strength=guidance_strength,
|
|
guidance_interval=guidance_interval,
|
|
progress=progress,
|
|
)
|
|
occ = decoder(latent)
|
|
mx.eval(occ)
|
|
return occ, latent
|
|
|
|
|
|
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
|
|
|
|
o = np.asarray(occ)[:, 0] # [B, D, H, W]
|
|
idx = np.argwhere(o > threshold).astype(np.int32)
|
|
return mx.array(idx)
|