Spaces:
Sleeping
Sleeping
| """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, | |
| ) |