pixal3d_mrp_mlx/pixal3d_mlx/pipeline.py

"""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

# 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,
    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,
    guidance_rescale: float = 0.0,
    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,
        guidance_rescale=guidance_rescale,
        progress=progress,
    )
    occ = decoder(latent)
    mx.eval(occ)
    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

    o = np.asarray(occ)[:, 0]  # [B, D, H, W]
    idx = np.argwhere(o > threshold).astype(np.int32)
    return mx.array(idx)