pixal3d_mrp_mlx/pixal3d_mlx/dino.py

"""DINOv3 image encoder, deliberately left in torch.

This is a stock ViT-L/16 run ONCE per image, outside the 25-step denoising loop, so it
is not worth porting to MLX: torch on MPS runs it fine and `transformers` gives exact
parity with upstream for free. The effort belongs in the 30-block DiT loops, which are
already ported and verified. The output is handed over as mx.arrays so nothing
downstream of here touches torch.

`facebook/dinov3-*` is gated; upstream itself falls back to the ungated `camenduru`
mirror (there is a HACK in pixal3d/pipelines/trellis2_image_to_3d.py doing exactly this
rewrite), so the mirror is the default here too.
"""

from __future__ import annotations

from typing import Tuple

import mlx.core as mx
import numpy as np

MODEL_NAME = "camenduru/dinov3-vitl16-pretrain-lvd1689m"

# ImageNet statistics — upstream's only transform (it assumes the image is pre-resized).
_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 3, 1, 1)
_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 3, 1, 1)


def _layers(model):
    """The transformer block list, across transformers versions.

    Upstream pins transformers==4.57.3 where `DINOv3ViTModel.layer` is a ModuleList on
    the model itself. In 5.x it moved to `model.model.layer`. Both are the same 24
    blocks; only the attribute path changed.
    """
    for owner in (model, getattr(model, "model", None)):
        if owner is None:
            continue
        for attr in ("layer", "layers"):
            found = getattr(owner, attr, None)
            if found is not None and hasattr(found, "__len__"):
                return found
    raise AttributeError("could not locate the DINOv3 transformer block list")


class DinoV3Encoder:
    """Wraps DINOv3 and returns (global tokens, patch feature map) as mx.arrays."""

    def __init__(self, model_name: str = MODEL_NAME, image_size: int = 512,
                 device: str | None = None):
        import torch
        from transformers import DINOv3ViTModel

        if device is None:
            device = "mps" if torch.backends.mps.is_available() else "cpu"
        self.device = device
        self.image_size = image_size

        self.model = DINOv3ViTModel.from_pretrained(model_name).eval().to(device)
        self.model.requires_grad_(False)
        self.layers = _layers(self.model)

        self.patch_size = self.model.config.patch_size
        self.patch_number = image_size // self.patch_size
        self.embed_dim = self.model.config.hidden_size
        self.num_register_tokens = getattr(self.model.config, "num_register_tokens", 4)

    def _hidden_states(self, image):
        """Upstream's `extract_features`, reproduced exactly.

        The final `layer_norm` is PARAMETERLESS — it is not `self.model.norm`, which
        carries weights. Substituting the model's own norm here changes the scale of
        every conditioning vector, and because it has no checkpoint trace it is
        precisely the class of bug that weight-key matching cannot catch (see CLAUDE.md
        — the same trap cost us a 200x-too-large ss_flow output).
        """
        import torch
        import torch.nn.functional as F

        image = image.to(self.model.embeddings.patch_embeddings.weight.dtype)
        hidden_states = self.model.embeddings(image, bool_masked_pos=None)
        position_embeddings = self.model.rope_embeddings(image)
        for layer_module in self.layers:
            hidden_states = layer_module(hidden_states, position_embeddings=position_embeddings)
            if isinstance(hidden_states, (tuple, list)):
                hidden_states = hidden_states[0]
        return F.layer_norm(hidden_states, hidden_states.shape[-1:])

    def __call__(self, image_chw01: np.ndarray) -> Tuple[mx.array, mx.array]:
        """[B,3,H,W] float32 in [0,1] -> (global [B,1+regs,D], patches [B,h,w,D]).

        `global` is only the CLS token plus the register tokens — the patch tokens do
        NOT go into it. They go into the proj branch, which is the whole point of the
        architecture: image detail reaches the DiT through back-projection, not
        through cross-attention.
        """
        import torch

        arr = (np.asarray(image_chw01, dtype=np.float32) - _MEAN) / _STD
        with torch.no_grad():
            z = self._hidden_states(torch.from_numpy(arr).to(self.device))
            z = z.float().cpu().numpy()

        b, _, d = z.shape
        n_reg = self.num_register_tokens
        z_global = z[:, : 1 + n_reg]                       # CLS + registers
        z_patch = z[:, 1 + n_reg:]                          # spatial tokens
        z_patch = z_patch.reshape(b, self.patch_number, self.patch_number, d)
        return mx.array(z_global), mx.array(z_patch)