Spaces:
Sleeping
Sleeping
File size: 5,151 Bytes
2267636 ec1d824 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 2267636 ec46c56 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """DINOv3 feature extraction via HuggingFace, handling CLS and register tokens."""
import os
import torch
import torch.nn as nn
from transformers import DINOv3ViTConfig, DINOv3ViTModel, DINOv3ViTImageProcessorFast
class DINOv3HFExtractor(nn.Module):
"""
Extracts intermediate features from DINOv3 via HuggingFace transformers.
Builds the model from config (no gated download required); weights are loaded
from the MMDiff checkpoint which bundles the DINOv3 backbone.
Returns 4 feature maps of shape [B, C_dino, H//16, W//16] from selected layers.
Input images must be [B, 3, H, W] in [0, 1] range.
"""
def __init__(self, repo_id="facebook/dinov3-vitb16-pretrain-lvd1689m",
take_last=None, take_indices=None, trainable=False,
hidden_size=768, num_hidden_layers=12, num_attention_heads=12,
intermediate_size=3072, patch_size=16, image_size=512,
num_register_tokens=4):
super().__init__()
# Build image processor from default config (no gated download needed)
self.proc = DINOv3ViTImageProcessorFast()
# Disable resizing/cropping so native resolution maps to patches
for k in ("do_resize", "do_center_crop"):
if hasattr(self.proc, k):
setattr(self.proc, k, False)
# Build model from config (random weights; real weights loaded from checkpoint)
config = DINOv3ViTConfig(
hidden_size=hidden_size,
num_hidden_layers=num_hidden_layers,
num_attention_heads=num_attention_heads,
intermediate_size=intermediate_size,
patch_size=patch_size,
image_size=image_size,
num_register_tokens=num_register_tokens,
hidden_act="gelu",
)
self.model = DINOv3ViTModel(config)
self.model.config.output_hidden_states = True
self._frozen = not trainable
if self._frozen:
self.model.eval()
for p in self.model.parameters():
p.requires_grad = False
else:
self.model.train()
for p in self.model.parameters():
p.requires_grad = True
# ImageNet normalization stats
mean = torch.tensor(self.proc.image_mean).view(1, 3, 1, 1)
std = torch.tensor(self.proc.image_std).view(1, 3, 1, 1)
self.register_buffer("mean", mean, persistent=False)
self.register_buffer("std", std, persistent=False)
if take_indices is not None:
self.take_indices = take_indices
self.take_last = None
else:
self.take_last = take_last if take_last is not None else 4
self.take_indices = None
self.patch_size = getattr(self.model.config, "patch_size", 16)
self.num_register_tokens = getattr(self.model.config, "num_register_tokens", 0)
hidden_size = getattr(self.model.config, "hidden_size", 768)
layers = self.take_indices if self.take_indices is not None else f"last {self.take_last}"
trainable_str = "trainable" if not self._frozen else "frozen"
print(f"[DINOv3] Built from config: dim={hidden_size}, patch={self.patch_size}, "
f"layers={layers} ({trainable_str})")
def train(self, mode: bool = True):
self.training = mode
if self._frozen:
self.model.eval()
else:
self.model.train(mode)
return self
def forward(self, images_512: torch.Tensor):
with torch.set_grad_enabled(not self._frozen):
return self._forward(images_512)
def _forward(self, images_512: torch.Tensor):
x = (images_512 - self.mean) / self.std
out = self.model(pixel_values=x, output_hidden_states=True)
hidden_states = out.hidden_states
B, _, H, W = images_512.shape
H_patches = H // self.patch_size
W_patches = W // self.patch_size
P = H_patches * W_patches
R = self.num_register_tokens
maps = []
if self.take_indices is not None:
for idx in self.take_indices:
hidden = hidden_states[idx]
spatial = hidden[:, 1:1+P, :]
C = spatial.shape[-1]
spatial_map = spatial.transpose(1, 2).reshape(B, C, H_patches, W_patches).contiguous()
maps.append(spatial_map)
else:
for hidden in hidden_states[-self.take_last:]:
spatial = hidden[:, 1:1+P, :]
C = spatial.shape[-1]
spatial_map = spatial.transpose(1, 2).reshape(B, C, H_patches, W_patches).contiguous()
maps.append(spatial_map)
return maps
def create_dinov3_hf_extractor(repo_id="facebook/dinov3-vitb16-pretrain-lvd1689m", take_last=None, take_indices=None, trainable=False):
"""
Factory for DINOv3HFExtractor (frozen in eval mode unless trainable=True).
Builds from config — no gated download required.
"""
return DINOv3HFExtractor(
repo_id=repo_id, take_last=take_last, take_indices=take_indices, trainable=trainable,
) |