Upload model.py with huggingface_hub

model.py

@@ -43,17 +43,49 @@ class AnyThermalSegmentationModel(PreTrainedModel):
         super().__init__(config)
         self.backbone = Dinov2Model(config)
         
-        # Wrapped in .model to match checkpoint "head.model.3.weight"
+        # Head definition matches your NonlinearHead64
         self.head = nn.Module()
         self.head.model = nn.Sequential(
             nn.Conv2d(config.hidden_size, 64, kernel_size=3, padding=1),
             nn.ReLU(inplace=True),
-            nn.Dropout2d(p=0.0), # Index 2 (Placeholder for structure)
+            nn.Dropout2d(p=0.0), 
             nn.Conv2d(64, config.num_labels, kernel_size=1)
         )
+        
+        # Define Normalization constants as buffers so they move to GPU automatically
+        self.register_buffer("norm_mean", torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1))
+        self.register_buffer("norm_std", torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1))
+        
         self.post_init()
 
-    def forward(self, pixel_values, **kwargs):
+    def preprocess_input(self, x):
+        """
+        Replicates preprocess_dinov2: 
+        1. Resize to nearest multiple of 14
+        2. Normalize with ViT stats
+        """
+        B, C, H, W = x.shape
+        patch_size = 14
+        
+        # 1. Dynamic Resize (Snap to grid)
+        new_H = (H // patch_size) * patch_size
+        new_W = (W // patch_size) * patch_size
+        
+        if new_H != H or new_W != W:
+            x = F.interpolate(x, size=(new_H, new_W), mode='bilinear', align_corners=False)
+            
+        # 2. Normalize
+
+        if x.max() > 1.0: x = x / 255.0
+        
+        x = (x - self.norm_mean) / self.norm_std
+        return x
+
+    def forward(self, pixel_values, labels=None, **kwargs):
+        # --- APPLY PREPROCESSING HERE ---
+        pixel_values = self.preprocess_input(pixel_values)
+        # --------------------------------
+        
         outputs = self.backbone(pixel_values, **kwargs)
         features = outputs.last_hidden_state[:, 1:, :] 
         B, L, C = features.shape
@@ -61,7 +93,17 @@ class AnyThermalSegmentationModel(PreTrainedModel):
         features = features.permute(0, 2, 1).reshape(B, C, H, W)
         
         logits = self.head.model(features)
-        return F.interpolate(logits, scale_factor=14, mode='bilinear', align_corners=False)
+        
+        # Upscale back to input size
+        logits = F.interpolate(logits, size=pixel_values.shape[-2:], mode='bilinear', align_corners=False)
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits, labels)
+            return {"loss": loss, "logits": logits}
+            
+        return logits
 
 
 # =============================================================================