Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,42 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, Dinov2Model, Dinov2Config
+
+class AnyThermalConfig(Dinov2Config):
+    model_type = "anythermal"
+
+class AnyThermalSegmentationModel(PreTrainedModel):
+    config_class = AnyThermalConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        # Initialize the backbone with DINOv2 architecture [cite: 100, 104]
+        self.backbone = Dinov2Model(config)
+        
+        # Build the segmentation head based on your non-linear-64 architecture [cite: 130]
+        # num_labels is defined in the subfolder config.json
+        self.head = nn.Sequential(
+            nn.Conv2d(config.hidden_size, 64, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, config.num_labels, kernel_size=1)
+        )
+        self.post_init()
+
+    def forward(self, pixel_values, **kwargs):
+        # Extract features [cite: 109, 131]
+        outputs = self.backbone(pixel_values, **kwargs)
+        
+        # Use patch tokens (skipping CLS token at 0) for segmentation [cite: 109, 131]
+        features = outputs.last_hidden_state[:, 1:, :] 
+        B, L, C = features.shape
+        H = W = int(L**0.5)
+        features = features.permute(0, 2, 1).reshape(B, C, H, W)
+        
+        logits = self.head(features)
+        
+        # Upscale to original resolution (14x) [cite: 131]
+        return F.interpolate(logits, scale_factor=14, mode='bilinear', align_corners=False)
+
+# Register for AutoModel discovery
+AnyThermalSegmentationModel.register_for_auto_class("AutoModel")