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

import time
from typing import Callable, Optional

import mlx.core as mx
import numpy as np

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 gather_proj_at_coords(cond: dict, coords: mx.array, grid_resolution: int) -> dict:
    """Reduce the dense [B, R^3, C] proj grid to just the occupied voxels.

    The structure stage conditions on every voxel of the lattice; the SLAT stage is
    sparse and only ever touches the occupied set, so the proj features are gathered
    at those coordinates. Coordinate order is (batch, x, y, z), matching upstream's
    `encode_image_proj`.

    `global` is not indexed — it is cross-attended — but it IS reshaped, because the
    sparse attention path takes a flat [M, C] stack with an explicit layout, whereas
    the dense structure stage takes [B, T, C]. Handing the dense shape to the sparse
    blocks fails inside `to_kv`'s reshape rather than anywhere informative.
    """
    r = grid_resolution
    proj = cond["proj"]
    b_size = proj.shape[0]
    if b_size != 1:
        raise NotImplementedError(
            "sparse conditioning assumes batch 1; a real batch needs an explicit "
            "context layout rather than the default single slice"
        )
    proj = proj.reshape(b_size, r, r, r, -1)
    b, x, y, z = (coords[:, i] for i in range(4))
    flat = ((b * r + x) * r + y) * r + z
    gathered = proj.reshape(-1, proj.shape[-1])[flat]
    g = cond["global"]
    return {"global": g.reshape(-1, g.shape[-1]), "proj": gathered}


def run_slat_stage(
    flow_model,
    cond: dict,
    coords: mx.array,
    neg_cond=None,
    normalization: dict | None = None,
    seed: int = 0,
    progress: Optional[Callable[[int, int], None]] = None,
    **overrides,
):
    """Sparse latent sampling seeded by the occupancy set. Returns a SparseTensor.

    The denormalisation at the end is NOT cosmetic: the flows are trained on
    standardised latents, and the decoders expect raw ones. Skipping it yields a mesh
    that decodes without error and is quietly the wrong scale.

    One stage of the cascade. `image_to_mesh` composes two of these — see there for the
    LR -> refine -> HR sequence upstream actually ships.
    """
    from trellis_sparse_mlx import SparseTensor

    params = {**SHAPE_SLAT_PARAMS, **overrides}
    mx.random.seed(seed)
    noise = SparseTensor(
        mx.random.normal((coords.shape[0], flow_model.in_channels)), coords
    )
    slat = FlowEulerSampler().sample(
        flow_model, noise, cond=cond, neg_cond=neg_cond, progress=progress, **params
    )
    if normalization:
        mean = mx.array(np.asarray(normalization["mean"], dtype=np.float32))[None]
        std = mx.array(np.asarray(normalization["std"], dtype=np.float32))[None]
        slat = slat.replace(slat.feats * std + mean)
    mx.eval(slat.feats)
    return slat


def run_tex_stage(
    flow_model,
    tex_decoder,
    cond: dict,
    shape_slat,
    subs,
    neg_cond=None,
    shape_normalization: dict | None = None,
    tex_normalization: dict | None = None,
    seed: int = 0,
    progress: Optional[Callable[[int, int], None]] = None,
    **overrides,
):
    """Texture SLAT + decode -> PBR voxels (base_color / metallic / roughness / alpha).

    Two things here are easy to get wrong and both fail silently rather than loudly:

    1. `shape_slat` arrives DENORMALISED (the shape stage un-standardises it for the
       decoder), but the texture flow was trained against the standardised form, so it
       is re-normalised before being used as `concat_cond`. Feeding the denormalised
       latent gives a plausible mesh with wrong colours.
    2. The texture decoder has `pred_subdiv=False` — it cannot invent its own
       subdivisions and must be handed the shape decoder's `subs` as guides, so the
       texture voxels land on exactly the geometry that was built.

    The decoder's tanh-ish output is mapped `* 0.5 + 0.5` into [0,1], which is the
    range o_voxel's baker expects.
    """
    from trellis_sparse_mlx import SparseTensor

    if shape_normalization:
        mean = mx.array(np.asarray(shape_normalization["mean"], dtype=np.float32))[None]
        std = mx.array(np.asarray(shape_normalization["std"], dtype=np.float32))[None]
        shape_slat = shape_slat.replace((shape_slat.feats - mean) / std)

    params = {**TEX_SLAT_PARAMS, **overrides}
    mx.random.seed(seed)
    noise_channels = flow_model.in_channels - shape_slat.feats.shape[1]
    noise = shape_slat.replace(
        mx.random.normal((shape_slat.coords.shape[0], noise_channels))
    )

    sampler = FlowEulerSampler(concat_cond=shape_slat)
    slat = sampler.sample(flow_model, noise, cond=cond, neg_cond=neg_cond,
                          progress=progress, **params)

    if tex_normalization:
        mean = mx.array(np.asarray(tex_normalization["mean"], dtype=np.float32))[None]
        std = mx.array(np.asarray(tex_normalization["std"], dtype=np.float32))[None]
        slat = slat.replace(slat.feats * std + mean)

    voxels = tex_decoder(slat, guide_subs=subs)
    voxels = voxels.replace(voxels.feats * 0.5 + 0.5)
    mx.eval(voxels.feats)
    return voxels


def refine_coords(shape_decoder, lr_slat, hr_resolution: int = 1024,
                  lr_resolution: int = 512, max_num_tokens: int = 49152):
    """Low-res SLAT -> refined high-res coordinate set. Stage 3a of the cascade.

    The low-res latent is pushed four stages into the shape decoder purely so its
    predicted subdivisions grow the occupied set; those coords (now at 16x, i.e. 512)
    are then quantised down to the high-res flow's own grid (1024 // 16 = 64) and
    deduplicated. This is the step whose absence leaves geometry outside the silhouette.

    Upstream backs the resolution off in 128 steps while the token count exceeds
    `max_num_tokens`, so a dense object degrades gracefully instead of exploding.
    Returns (coords, hr_resolution, grid_res).
    """
    import numpy as _np

    hr_coords = _np.asarray(shape_decoder.upsample(lr_slat, upsample_times=4))
    res = hr_resolution
    while True:
        grid_res = res // 16
        q = _np.concatenate([
            hr_coords[:, :1],
            _np.rint((hr_coords[:, 1:] + 0.5) / lr_resolution * (grid_res - 1)).astype(_np.int32),
        ], axis=1)
        uniq = _np.unique(q, axis=0)
        if uniq.shape[0] < max_num_tokens or res == 1024:
            break
        res -= 128
    return mx.array(uniq.astype(_np.int32)), res, res // 16


def image_to_mesh(
    image_path: str,
    models: dict,
    camera_angle_x: float = DEFAULT_FOV,
    mesh_scale: float = 1.0,
    seed: int = 0,
    normalization: dict | None = None,
    tex_normalization: dict | None = None,
    texture: bool = False,
    log=print,
):
    """The full shipped geometry cascade: image -> (vertices, faces).

    `models` supplies the five pieces, so a long-lived server can load them ONCE:
    `ss_flow`, `ss_dec`, `slat_512`, `slat_1024`, `shape_dec`, plus the conditioners
    `cond_512` and `cond_1024`.

    The sequence, and why each step is where it is:

      1. structure  -> 64^3 occupancy, MAX-POOLED DOWN to 32^3. The cascade starts low.
      2. LR SLAT    -> sparse latents at 32^3, conditioned by the shape_512 extractor
      3. refine     -> the LR latent is pushed four decoder stages so its predicted
                       subdivisions grow the occupied set; those coords quantise to 64^3
      4. HR SLAT    -> sparse latents at the refined coords, shape_1024 extractor
      5. decode     -> Flexible Dual Grid at 1024^3 -> triangles

    Returns (vertices, faces, info). Vertices are in the VOXEL GRID frame; use
    `mesh.to_camera_frame` before comparing against the source image.
    """
    from .cond import load_image
    from .mesh import fdg_to_mesh, output_resolution

    t0 = time.time()
    occ, _, _ = image_to_occupancy(image_path, models["ss_flow"], models["ss_dec"],
                                   camera_angle_x=camera_angle_x,
                                   mesh_scale=mesh_scale, seed=seed)
    coords = occupied_coords_at(occ, 32)
    log(f"  structure  {coords.shape[0]:6} voxels @32^3   {time.time() - t0:6.1f}s")

    t = time.time()
    c512 = models["cond_512"]
    img = load_image(image_path, c512.image_size)
    cond, uncond = c512(img, camera_angle_x, mesh_scale=mesh_scale)
    lr_slat = run_slat_stage(
        models["slat_512"], gather_proj_at_coords(cond, coords, 32), coords,
        neg_cond=gather_proj_at_coords(uncond, coords, 32),
        normalization=normalization, seed=seed)
    log(f"  LR SLAT    {lr_slat.feats.shape[0]:6} x {lr_slat.feats.shape[1]}          "
        f"{time.time() - t:6.1f}s")

    t = time.time()
    hr_coords, hr_res, grid_res = refine_coords(models["shape_dec"], lr_slat)
    log(f"  refine     {hr_coords.shape[0]:6} coords @{grid_res}^3 (res {hr_res})  "
        f"{time.time() - t:6.1f}s")

    t = time.time()
    c1024 = models["cond_1024"]
    img = load_image(image_path, c1024.image_size)
    cond, uncond = c1024(img, camera_angle_x, mesh_scale=mesh_scale,
                         grid_resolution=grid_res)
    hr_slat = run_slat_stage(
        models["slat_1024"], gather_proj_at_coords(cond, hr_coords, grid_res), hr_coords,
        neg_cond=gather_proj_at_coords(uncond, hr_coords, grid_res),
        normalization=normalization, seed=seed)
    log(f"  HR SLAT    {hr_slat.feats.shape[0]:6} x {hr_slat.feats.shape[1]}          "
        f"{time.time() - t:6.1f}s")

    t = time.time()
    out, subs = models["shape_dec"](hr_slat, return_subs=True)
    mx.eval(out.feats)
    res_out = output_resolution(out)
    v, f = fdg_to_mesh(out, res_out)
    log(f"  mesh       {v.shape[0]:6} verts, {f.shape[0]} faces @{res_out}^3  "
        f"{time.time() - t:6.1f}s")

    info = {"hr_resolution": hr_res, "grid_resolution": grid_res,
            "output_resolution": res_out, "num_tokens": int(hr_coords.shape[0]),
            "peak_gb": round(mx.get_peak_memory() / 2 ** 30, 1)}
    extra = {"subs": subs, "hr_slat": hr_slat}

    if texture:
        t = time.time()
        ctex = models["cond_tex"]
        img = load_image(image_path, ctex.image_size)
        cond, uncond = ctex(img, camera_angle_x, mesh_scale=mesh_scale,
                            grid_resolution=grid_res)
        voxels = run_tex_stage(
            models["slat_tex"], models["tex_dec"],
            gather_proj_at_coords(cond, hr_coords, grid_res), hr_slat, subs,
            neg_cond=gather_proj_at_coords(uncond, hr_coords, grid_res),
            shape_normalization=normalization, tex_normalization=tex_normalization,
            seed=seed)
        log(f"  texture    {voxels.feats.shape[0]:6} PBR voxels x "
            f"{voxels.feats.shape[1]}      {time.time() - t:6.1f}s")
        extra["tex_voxels"] = voxels
        info["textured"] = True

    info["seconds"] = round(time.time() - t0, 1)
    info["peak_gb"] = round(mx.get_peak_memory() / 2 ** 30, 1)
    return v, f, {**info, **extra}


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_at(occ: mx.array, resolution: int, threshold: float = 0.0) -> mx.array:
    """Occupancy logits -> int32 [N,4] (batch, x, y, z) coords at `resolution`.

    ss_dec always decodes to 64^3, but the cascade STARTS AT 32^3 — upstream max-pools
    the boolean grid down by the ratio (`max_pool3d(..., ratio) > 0.5`, i.e. a voxel
    survives if ANY of its eight children was occupied). Feeding the raw 64^3 set into
    the high-res flow instead skips the low-res refinement entirely and leaves a halo
    of geometry outside the true silhouette.
    """
    import numpy as _np

    o = _np.asarray(occ)[:, 0] > threshold          # [B,D,H,W] bool
    d = o.shape[1]
    if resolution != d:
        ratio = d // resolution
        if d % resolution:
            raise ValueError(f"{d} is not an integer multiple of {resolution}")
        o = o.reshape(o.shape[0], resolution, ratio, resolution, ratio, resolution, ratio)
        o = o.any(axis=(2, 4, 6))                   # max-pool over the boolean grid
    return mx.array(_np.argwhere(o).astype(_np.int32))


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)