Upload models/dinov3_hf_extractor.py with huggingface_hub

models/dinov3_hf_extractor.py

@@ -3,39 +3,49 @@
 import os
 import torch
 import torch.nn as nn
-from transformers import AutoImageProcessor, AutoModel
+from transformers import DINOv3ViTConfig, DINOv3ViTModel, DINOv3ViTImageProcessorFast
 
 
 class DINOv3HFExtractor(nn.Module):
     """
     Extracts intermediate features from DINOv3 via HuggingFace transformers.
-    
-    Properly handles DINOv3 register tokens:
-    - Output shape: [B, 1 + P + R, C] where:
-      - 1 = CLS token
-      - P = spatial patches (1024 for 512x512 with 16x16 patches)
-      - R = register tokens (typically 4 for DINOv3)
-    
-    Returns 4 feature maps of shape [B, C_dino, 32, 32] from last 4 layers.
-    Input images must be [B, 3, 512, 512] in [0, 1] range.
+
+    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):
+
+    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__()
 
-        token = os.environ.get("HF_TOKEN")
-        self.proc = AutoImageProcessor.from_pretrained(repo_id, token=token)
-        
-        # Disable resizing/cropping so native 512x512 maps to 32x32 patches
+        # 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)
-        
-        self.model = AutoModel.from_pretrained(repo_id, token=token)
+
+        # 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
-        
-        # trainable=True is used by the `dino_only` ablation (fine-tune the backbone);
-        # otherwise the backbone is frozen and kept in eval mode.
+
         self._frozen = not trainable
         if self._frozen:
             self.model.eval()
@@ -45,93 +55,77 @@ class DINOv3HFExtractor(nn.Module):
             self.model.train()
             for p in self.model.parameters():
                 p.requires_grad = True
-        
-        # ImageNet normalization stats: mean [0.485, 0.456, 0.406], std [0.229, 0.224, 0.225]
+
+        # 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] {repo_id.split('/')[-1]}: dim={hidden_size}, patch={self.patch_size}, "
+        print(f"[DINOv3] Built from config: dim={hidden_size}, patch={self.patch_size}, "
               f"layers={layers} ({trainable_str})")
-    
+
     def train(self, mode: bool = True):
-        """Keep a frozen backbone in eval mode; otherwise follow `mode`."""
         self.training = mode
         if self._frozen:
             self.model.eval()
         else:
             self.model.train(mode)
         return self
-    
+
     def forward(self, images_512: torch.Tensor):
-        """Extract DINOv3 features from [B, 3, H, W] images in [0, 1]; returns maps [B, C, H//16, W//16]."""
-        # Disable grad only when the backbone is frozen; otherwise allow fine-tuning
         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  # Tuple: [emb, blk1, ..., blkL]
-        
+        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  # total spatial patches
+        P = H_patches * W_patches
         R = self.num_register_tokens
-        
-        # Each hidden state is [B, 1 + P + R, C], laid out as [CLS][P spatial patches][R 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, :]  # [B, P, C], drop CLS and register tokens
+                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, :]  # [B, P, C], drop CLS and register tokens
+                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).
-    repo_id options: vits16 (384-dim), vitb16 (768-dim, recommended), vitl16 (1024-dim).
+    Builds from config — no gated download required.
     """
-    return DINOv3HFExtractor(repo_id=repo_id, take_last=take_last, take_indices=take_indices, trainable=trainable)
-
-
-if __name__ == "__main__":
-    extractor = create_dinov3_hf_extractor().cuda()
-    
-    dummy_input = torch.randn(2, 3, 512, 512).cuda()
-    print(f"\nInput: {dummy_input.shape}")
-    
-    features = extractor(dummy_input)
-    
-    print(f"\nExtracted {len(features)} feature maps:")
-    for i, feat in enumerate(features):
-        print(f"  Layer {i}: {feat.shape}")
-
+    return DINOv3HFExtractor(
+        repo_id=repo_id, take_last=take_last, take_indices=take_indices, trainable=trainable,
+    )