Upload model.py with huggingface_hub

model.py

@@ -38,5 +38,75 @@ class AnyThermalSegmentationModel(PreTrainedModel):
         # Upscale to original resolution (14x) [cite: 131]
         return F.interpolate(logits, scale_factor=14, mode='bilinear', align_corners=False)
 
+
+# 1. Custom Config to handle SALAD parameters
+class AnyThermalVPRConfig(Dinov2Config):
+    model_type = "anythermal_vpr"
+    def __init__(self, num_clusters=64, cluster_dim=128, token_dim=256, **kwargs):
+        super().__init__(**kwargs)
+        self.num_clusters = num_clusters
+        self.cluster_dim = cluster_dim
+        self.token_dim = token_dim
+
+# 2. SALAD Aggregator (Logic from salad.py)
+class SALADHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.num_channels = config.hidden_size
+        self.num_clusters = config.num_clusters
+        self.cluster_dim = config.cluster_dim
+        self.token_dim = config.token_dim
+
+        self.token_features = nn.Sequential(
+            nn.Linear(self.num_channels, 512),
+            nn.ReLU(),
+            nn.Linear(512, self.token_dim)
+        )
+        self.cluster_features = nn.Sequential(
+            nn.Conv2d(self.num_channels, 512, 1),
+            nn.ReLU(),
+            nn.Conv2d(512, self.cluster_dim, 1)
+        )
+        self.score = nn.Sequential(
+            nn.Conv2d(self.num_channels, 512, 1),
+            nn.ReLU(),
+            nn.Conv2d(512, self.num_clusters, 1),
+        )
+
+    def forward(self, x_tuple):
+        x, t = x_tuple # patch features [B, C, H/14, W/14], cls token [B, C]
+        f = self.cluster_features(x).flatten(2)
+        p = F.softmax(self.score(x).flatten(2), dim=1) # Simplified Sinkhorn for inference
+        t = self.token_features(t)
+
+        vlad = (f.unsqueeze(2).repeat(1, 1, self.num_clusters, 1) * p.unsqueeze(1)).sum(dim=-1)
+        vlad = F.normalize(vlad, p=2, dim=1).flatten(1)
+        
+        combined = torch.cat([F.normalize(t, p=2, dim=-1), vlad], dim=-1)
+        return F.normalize(combined, p=2, dim=-1)
+
+# 3. Final VPR Model
+class AnyThermalVPRModel(PreTrainedModel):
+    config_class = AnyThermalVPRConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.backbone = Dinov2Model(config)
+        self.vpr_head = SALADHead(config)
+        self.post_init()
+
+    def forward(self, pixel_values, **kwargs):
+        outputs = self.backbone(pixel_values, **kwargs)
+        # Prepare inputs for SALAD
+        patch_tokens = outputs.last_hidden_state[:, 1:, :].permute(0, 2, 1)
+        B, C, L = patch_tokens.shape
+        H = W = int(L**0.5)
+        patch_tokens = patch_tokens.reshape(B, C, H, W)
+        cls_token = outputs.last_hidden_state[:, 0, :]
+        
+        # Global descriptor
+        return self.vpr_head((patch_tokens, cls_token))
+
+AnyThermalVPRModel.register_for_auto_class("AutoModel")
 # Register for AutoModel discovery
 AnyThermalSegmentationModel.register_for_auto_class("AutoModel")