350 lines
14 KiB
Python
350 lines
14 KiB
Python
"""
|
|
DINOv3 ViT feature extractor in MLX.
|
|
Uses HuggingFace transformers weights but runs purely in MLX.
|
|
"""
|
|
import math
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
import numpy as np
|
|
from PIL import Image
|
|
from typing import Optional
|
|
|
|
|
|
class MlxDINOv3PatchEmbed(nn.Module):
|
|
"""Patch embedding: Conv2d(3, dim, kernel=16, stride=16) + CLS + registers."""
|
|
|
|
def __init__(self, dim: int = 1024, patch_size: int = 16, num_register_tokens: int = 4):
|
|
super().__init__()
|
|
self.patch_size = patch_size
|
|
self.num_register_tokens = num_register_tokens
|
|
self.weight = mx.zeros((dim, patch_size, patch_size, 3))
|
|
self.bias = mx.zeros((dim,))
|
|
self.cls_token = mx.zeros((1, 1, dim))
|
|
self.register_tokens = mx.zeros((1, num_register_tokens, dim))
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
"""x: (B, H, W, 3) -> (B, num_patches + 1 + num_reg, dim)"""
|
|
B, H, W, C = x.shape
|
|
P = self.patch_size
|
|
nH, nW = H // P, W // P
|
|
|
|
# Manual convolution via reshape + matmul (stride == kernel_size)
|
|
x = x.reshape(B, nH, P, nW, P, C)
|
|
x = x.transpose(0, 1, 3, 2, 4, 5) # (B, nH, nW, P, P, C)
|
|
x = x.reshape(B, nH * nW, P * P * C)
|
|
|
|
w = self.weight.reshape(self.weight.shape[0], -1).T # (P*P*3, dim)
|
|
patches = x @ w + self.bias # (B, N, dim)
|
|
|
|
cls = mx.broadcast_to(self.cls_token, (B, 1, patches.shape[-1]))
|
|
reg = mx.broadcast_to(self.register_tokens, (B, self.num_register_tokens, patches.shape[-1]))
|
|
return mx.concatenate([cls, reg, patches], axis=1)
|
|
|
|
|
|
class MlxDINOv3TransformerBlock(nn.Module):
|
|
"""ViT transformer block with layer scaling."""
|
|
|
|
def __init__(self, dim: int = 1024, num_heads: int = 16, mlp_ratio: float = 4.0):
|
|
super().__init__()
|
|
self.norm1 = nn.LayerNorm(dim)
|
|
self.norm2 = nn.LayerNorm(dim)
|
|
self.attn = MlxDINOv3Attention(dim, num_heads)
|
|
hidden = int(dim * mlp_ratio)
|
|
self.mlp = MlxDINOv3MLP(dim, hidden)
|
|
# Layer scaling
|
|
self.layer_scale1 = mx.ones((dim,))
|
|
self.layer_scale2 = mx.ones((dim,))
|
|
|
|
def __call__(self, x: mx.array, rope_cos: mx.array = None, rope_sin: mx.array = None,
|
|
num_prefix_tokens: int = 5) -> mx.array:
|
|
h = self.norm1(x)
|
|
h = self.attn(h, rope_cos, rope_sin, num_prefix_tokens=num_prefix_tokens)
|
|
x = x + self.layer_scale1 * h
|
|
h = self.norm2(x)
|
|
h = self.mlp(h)
|
|
x = x + self.layer_scale2 * h
|
|
return x
|
|
|
|
|
|
class MlxDINOv3Attention(nn.Module):
|
|
"""Multi-head self-attention with RoPE (matching HF DINOv3ViTAttention).
|
|
|
|
Uses separate Q, K, V projections (not fused QKV) to match HF weight layout.
|
|
RoPE is applied only to patch tokens, skipping prefix (CLS + register) tokens.
|
|
"""
|
|
|
|
def __init__(self, dim: int, num_heads: int):
|
|
super().__init__()
|
|
self.num_heads = num_heads
|
|
self.head_dim = dim // num_heads
|
|
self.q_proj = nn.Linear(dim, dim, bias=True)
|
|
self.k_proj = nn.Linear(dim, dim, bias=False) # DINOv3: key has no bias
|
|
self.v_proj = nn.Linear(dim, dim, bias=True)
|
|
self.o_proj = nn.Linear(dim, dim, bias=True)
|
|
|
|
def __call__(self, x: mx.array, rope_cos: mx.array = None, rope_sin: mx.array = None,
|
|
num_prefix_tokens: int = 5) -> mx.array:
|
|
B, N, C = x.shape
|
|
H = self.num_heads
|
|
D = self.head_dim
|
|
|
|
q = self.q_proj(x).reshape(B, N, H, D).transpose(0, 2, 1, 3) # (B, H, N, D)
|
|
k = self.k_proj(x).reshape(B, N, H, D).transpose(0, 2, 1, 3)
|
|
v = self.v_proj(x).reshape(B, N, H, D).transpose(0, 2, 1, 3)
|
|
|
|
if rope_cos is not None:
|
|
# Apply RoPE only to patch tokens, skip prefix (CLS + registers)
|
|
num_patches = rope_cos.shape[-2]
|
|
q_prefix = q[:, :, :num_prefix_tokens]
|
|
k_prefix = k[:, :, :num_prefix_tokens]
|
|
q_patches = q[:, :, num_prefix_tokens:]
|
|
k_patches = k[:, :, num_prefix_tokens:]
|
|
|
|
q_patches = self._apply_rope(q_patches, rope_cos, rope_sin)
|
|
k_patches = self._apply_rope(k_patches, rope_cos, rope_sin)
|
|
|
|
q = mx.concatenate([q_prefix, q_patches], axis=2)
|
|
k = mx.concatenate([k_prefix, k_patches], axis=2)
|
|
|
|
scale = D ** -0.5
|
|
out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
|
|
out = out.transpose(0, 2, 1, 3).reshape(B, N, C)
|
|
return self.o_proj(out)
|
|
|
|
@staticmethod
|
|
def _apply_rope(x: mx.array, cos: mx.array, sin: mx.array) -> mx.array:
|
|
"""Apply RoPE via rotate_half: x*cos + rotate_half(x)*sin.
|
|
|
|
cos/sin: (1, N_patches, head_dim) — full head_dim, already tiled.
|
|
x: (B, H, N_patches, head_dim).
|
|
"""
|
|
# rotate_half: [-x2, x1]
|
|
half = x.shape[-1] // 2
|
|
x1 = x[..., :half]
|
|
x2 = x[..., half:]
|
|
x_rot = mx.concatenate([-x2, x1], axis=-1)
|
|
return x * cos + x_rot * sin
|
|
|
|
|
|
class MlxDINOv3MLP(nn.Module):
|
|
"""MLP with GELU activation."""
|
|
|
|
def __init__(self, dim: int, hidden: int):
|
|
super().__init__()
|
|
self.fc1 = nn.Linear(dim, hidden, bias=True)
|
|
self.fc2 = nn.Linear(hidden, dim, bias=True)
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
return self.fc2(nn.gelu(self.fc1(x)))
|
|
|
|
|
|
class MlxDINOv3FeatureExtractor(nn.Module):
|
|
"""
|
|
DINOv3 ViT-L/16 feature extractor in pure MLX.
|
|
24 layers, dim=1024, 16 heads, patch_size=16.
|
|
"""
|
|
|
|
def __init__(self, dim: int = 1024, num_heads: int = 16, num_layers: int = 24,
|
|
patch_size: int = 16, mlp_ratio: float = 4.0,
|
|
num_register_tokens: int = 4, rope_theta: float = 100.0):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.num_heads = num_heads
|
|
self.patch_size = patch_size
|
|
self.num_register_tokens = num_register_tokens
|
|
self.rope_theta = rope_theta
|
|
|
|
self.embeddings = MlxDINOv3PatchEmbed(dim, patch_size, num_register_tokens)
|
|
self.layers = [MlxDINOv3TransformerBlock(dim, num_heads, mlp_ratio) for _ in range(num_layers)]
|
|
self.norm = nn.LayerNorm(dim)
|
|
|
|
def _build_rope_2d(self, h: int, w: int) -> tuple:
|
|
"""Build 2D RoPE cos/sin matching HF DINOv3ViTRopePositionEmbedding.
|
|
|
|
Returns cos/sin for PATCH tokens only (prefix tokens get no RoPE).
|
|
Shape: (1, h*w, head_dim) — full head_dim, tiled 2x.
|
|
"""
|
|
head_dim = self.dim // self.num_heads
|
|
|
|
# inv_freq: 1/theta^(arange(0, 1, 4/head_dim)) — matches HF exactly
|
|
inv_freq = 1.0 / (self.rope_theta ** mx.arange(0, 1, 4 / head_dim, dtype=mx.float32))
|
|
# inv_freq shape: (head_dim/4,)
|
|
|
|
# Patch center coords normalized to [-1, +1] — matches HF get_patches_center_coordinates
|
|
coords_h = (mx.arange(h, dtype=mx.float32) + 0.5) / h
|
|
coords_w = (mx.arange(w, dtype=mx.float32) + 0.5) / w
|
|
grid_h, grid_w = mx.meshgrid(coords_h, coords_w, indexing='ij')
|
|
coords = mx.stack([grid_h.reshape(-1), grid_w.reshape(-1)], axis=-1) # (h*w, 2)
|
|
coords = 2.0 * coords - 1.0 # shift to [-1, +1]
|
|
|
|
# angles: 2π * coord * inv_freq, then flatten and tile 2x
|
|
angles = 2 * math.pi * coords[:, :, None] * inv_freq[None, None, :] # (h*w, 2, head_dim/4)
|
|
angles = angles.reshape(h * w, -1) # (h*w, head_dim/2)
|
|
# Tile to full head_dim (matching PT's angles.tile(2))
|
|
angles = mx.concatenate([angles, angles], axis=-1) # (h*w, head_dim)
|
|
|
|
cos = mx.cos(angles)[None, :, :] # (1, h*w, head_dim)
|
|
sin = mx.sin(angles)[None, :, :]
|
|
return cos, sin
|
|
|
|
def __call__(self, images: list) -> mx.array:
|
|
x = self._preprocess(images)
|
|
B, H, W, _ = x.shape
|
|
|
|
h = self.embeddings(x)
|
|
|
|
nH, nW = H // self.patch_size, W // self.patch_size
|
|
rope_cos, rope_sin = self._build_rope_2d(nH, nW)
|
|
|
|
# 1 CLS + num_register_tokens registers = prefix tokens (no RoPE)
|
|
num_prefix = 1 + self.num_register_tokens
|
|
|
|
for layer in self.layers:
|
|
h = layer(h, rope_cos, rope_sin, num_prefix_tokens=num_prefix)
|
|
|
|
# Match PT reference: pure layer_norm without learned params
|
|
# (PT's extract_features uses F.layer_norm, not the model's learned norm)
|
|
h = mx.fast.layer_norm(h, weight=None, bias=None, eps=1e-5)
|
|
return h
|
|
|
|
def _preprocess(self, images: list) -> mx.array:
|
|
"""Resize, normalize, convert to MLX array."""
|
|
mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
|
std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
|
|
|
processed = []
|
|
for img in images:
|
|
if isinstance(img, Image.Image):
|
|
size = max(img.size)
|
|
target = ((size + self.patch_size - 1) // self.patch_size) * self.patch_size
|
|
img = img.resize((target, target), Image.LANCZOS)
|
|
arr = np.array(img.convert('RGB')).astype(np.float32) / 255.0
|
|
else:
|
|
arr = np.array(img, dtype=np.float32)
|
|
arr = (arr - mean) / std
|
|
processed.append(arr)
|
|
|
|
return mx.array(np.stack(processed))
|
|
|
|
|
|
def load_dinov3_from_hf(
|
|
model_name: str = "facebook/dinov3-vitl16-pretrain-lvd1689m",
|
|
image_size: int = 512,
|
|
*,
|
|
revision: Optional[str] = None,
|
|
cache_dir: Optional[str] = None,
|
|
local_files_only: bool = False,
|
|
) -> MlxDINOv3FeatureExtractor:
|
|
"""Load DINOv3 weights from HuggingFace into MLX model."""
|
|
from huggingface_hub import hf_hub_download
|
|
import json
|
|
|
|
hub_kwargs = {
|
|
"revision": revision,
|
|
"cache_dir": cache_dir,
|
|
"local_files_only": local_files_only,
|
|
}
|
|
config_path = hf_hub_download(model_name, "config.json", **hub_kwargs)
|
|
weight_path = hf_hub_download(model_name, "model.safetensors", **hub_kwargs)
|
|
|
|
with open(config_path) as f:
|
|
config = json.load(f)
|
|
|
|
dim = config.get('hidden_size', 1024)
|
|
num_heads = config.get('num_attention_heads', 16)
|
|
num_layers = config.get('num_hidden_layers', 24)
|
|
patch_size = config.get('patch_size', 16)
|
|
mlp_ratio = config.get('mlp_ratio', 4.0)
|
|
num_register = config.get('num_register_tokens', 4)
|
|
|
|
model = MlxDINOv3FeatureExtractor(
|
|
dim=dim, num_heads=num_heads, num_layers=num_layers,
|
|
patch_size=patch_size, mlp_ratio=mlp_ratio,
|
|
num_register_tokens=num_register,
|
|
)
|
|
|
|
weights = mx.load(weight_path)
|
|
remapped = _remap_dinov3_weights(weights, num_layers, dim)
|
|
model.load_weights(list(remapped.items()))
|
|
|
|
return model
|
|
|
|
|
|
def _remap_dinov3_weights(weights: dict, num_layers: int, dim: int) -> dict:
|
|
"""
|
|
Map HuggingFace DINOv3 weight keys to our model structure.
|
|
|
|
HF format:
|
|
layer.N.attention.q_proj.weight, layer.N.attention.k_proj.weight (no bias),
|
|
layer.N.attention.v_proj.weight, layer.N.attention.o_proj.weight/bias,
|
|
layer.N.mlp.up_proj.weight/bias, layer.N.mlp.down_proj.weight/bias,
|
|
layer.N.norm1/norm2.weight/bias, layer.N.layer_scale1/2.lambda1
|
|
|
|
Our format:
|
|
layers.N.attn.qkv.weight/bias, layers.N.attn.proj.weight/bias,
|
|
layers.N.mlp.fc1.weight/bias, layers.N.mlp.fc2.weight/bias,
|
|
layers.N.norm1/norm2.weight/bias, layers.N.layer_scale1/2
|
|
"""
|
|
remapped = {}
|
|
|
|
# Embeddings
|
|
if 'embeddings.patch_embeddings.weight' in weights:
|
|
# HF: (dim, 3, P, P) -> our: (dim, P, P, 3)
|
|
remapped['embeddings.weight'] = weights['embeddings.patch_embeddings.weight'].transpose(0, 2, 3, 1)
|
|
if 'embeddings.patch_embeddings.bias' in weights:
|
|
remapped['embeddings.bias'] = weights['embeddings.patch_embeddings.bias']
|
|
if 'embeddings.cls_token' in weights:
|
|
remapped['embeddings.cls_token'] = weights['embeddings.cls_token']
|
|
if 'embeddings.register_tokens' in weights:
|
|
remapped['embeddings.register_tokens'] = weights['embeddings.register_tokens']
|
|
|
|
# Final norm
|
|
if 'norm.weight' in weights:
|
|
remapped['norm.weight'] = weights['norm.weight']
|
|
if 'norm.bias' in weights:
|
|
remapped['norm.bias'] = weights['norm.bias']
|
|
|
|
# Transformer layers — separate Q/K/V projections (matching HF layout)
|
|
for i in range(num_layers):
|
|
prefix_hf = f'layer.{i}'
|
|
prefix_mlx = f'layers.{i}'
|
|
|
|
# Q, K, V projections — keep separate (not fused)
|
|
for proj in ['q_proj', 'k_proj', 'v_proj']:
|
|
for suffix in ['weight', 'bias']:
|
|
key = f'{prefix_hf}.attention.{proj}.{suffix}'
|
|
if key in weights:
|
|
remapped[f'{prefix_mlx}.attn.{proj}.{suffix}'] = weights[key]
|
|
|
|
# Output projection
|
|
for suffix in ['weight', 'bias']:
|
|
key = f'{prefix_hf}.attention.o_proj.{suffix}'
|
|
if key in weights:
|
|
remapped[f'{prefix_mlx}.attn.o_proj.{suffix}'] = weights[key]
|
|
|
|
# Norms
|
|
for norm in ['norm1', 'norm2']:
|
|
for suffix in ['weight', 'bias']:
|
|
key = f'{prefix_hf}.{norm}.{suffix}'
|
|
if key in weights:
|
|
remapped[f'{prefix_mlx}.{norm}.{suffix}'] = weights[key]
|
|
|
|
# MLP: up_proj → fc1, down_proj → fc2
|
|
for suffix in ['weight', 'bias']:
|
|
key = f'{prefix_hf}.mlp.up_proj.{suffix}'
|
|
if key in weights:
|
|
remapped[f'{prefix_mlx}.mlp.fc1.{suffix}'] = weights[key]
|
|
key = f'{prefix_hf}.mlp.down_proj.{suffix}'
|
|
if key in weights:
|
|
remapped[f'{prefix_mlx}.mlp.fc2.{suffix}'] = weights[key]
|
|
|
|
# Layer scaling
|
|
key1 = f'{prefix_hf}.layer_scale1.lambda1'
|
|
if key1 in weights:
|
|
remapped[f'{prefix_mlx}.layer_scale1'] = weights[key1]
|
|
key2 = f'{prefix_hf}.layer_scale2.lambda1'
|
|
if key2 in weights:
|
|
remapped[f'{prefix_mlx}.layer_scale2'] = weights[key2]
|
|
|
|
return remapped
|