Upload model.py with huggingface_hub

model.py

@@ -107,6 +107,100 @@ class AnyThermalVPRModel(PreTrainedModel):
         # Global descriptor
         return self.vpr_head((patch_tokens, cls_token))
 
+class AnyThermalDepthConfig(Dinov2Config):
+    model_type = "anythermal_depth"
+    def __init__(self, features=256, **kwargs):
+        super().__init__(**kwargs)
+        self.features = features
+
+class ResidualConvUnit(nn.Module):
+    def __init__(self, features):
+        super().__init__()
+        self.conv1 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True)
+        self.conv2 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True)
+        self.relu = nn.ReLU(inplace=True)
+    def forward(self, x):
+        out = self.relu(x)
+        out = self.conv1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        return out + x
+
+class FeatureFusionBlock(nn.Module):
+    def __init__(self, features):
+        super().__init__()
+        self.resConfUnit1 = ResidualConvUnit(features)
+        self.resConfUnit2 = ResidualConvUnit(features)
+    def forward(self, *xs):
+        output = xs[0]
+        if len(xs) == 2:
+            output = output + self.resConfUnit1(xs[1])
+        output = self.resConfUnit2(output)
+        output = F.interpolate(output, scale_factor=2, mode="bilinear", align_corners=True)
+        return output
+
+class AnyThermalDepthModel(PreTrainedModel):
+    config_class = AnyThermalDepthConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.backbone = Dinov2Model(config)
+        features = config.features
+        
+        # Layers to match 'scratch' in blocks.py
+        self.scratch = nn.Module()
+        self.scratch.layer1_rn = nn.Conv2d(96, features, kernel_size=3, padding=1, bias=False)
+        self.scratch.layer2_rn = nn.Conv2d(192, features, kernel_size=3, padding=1, bias=False)
+        self.scratch.layer3_rn = nn.Conv2d(384, features, kernel_size=3, padding=1, bias=False)
+        self.scratch.layer4_rn = nn.Conv2d(768, features, kernel_size=3, padding=1, bias=False)
+        
+        # Post-processing from vit.py
+        self.act_postprocess1 = nn.Sequential(nn.Conv2d(768, 96, 1), nn.ConvTranspose2d(96, 96, 4, stride=4))
+        self.act_postprocess2 = nn.Sequential(nn.Conv2d(768, 192, 1), nn.ConvTranspose2d(192, 192, 2, stride=2))
+        self.act_postprocess3 = nn.Sequential(nn.Conv2d(768, 384, 1))
+        self.act_postprocess4 = nn.Sequential(nn.Conv2d(768, 768, 1), nn.Conv2d(768, 768, 3, stride=2, padding=1))
+        
+        # Fusion and output
+        self.refinenet4 = FeatureFusionBlock(features)
+        self.refinenet3 = FeatureFusionBlock(features)
+        self.refinenet2 = FeatureFusionBlock(features)
+        self.refinenet1 = FeatureFusionBlock(features)
+        self.output_conv = nn.Sequential(
+            nn.Conv2d(features, 128, kernel_size=3, padding=1),
+            nn.Upsample(scale_factor=1.75, mode="bilinear"), # Specific to Dinov2-ViT-B14
+            nn.Conv2d(128, 32, kernel_size=3, padding=1),
+            nn.ReLU(True),
+            nn.Conv2d(32, 1, kernel_size=1),
+            nn.ReLU(True)
+        )
+        self.post_init()
+
+    def forward(self, pixel_values):
+        # Extract features from layers 2, 5, 8, 11
+        outputs = self.backbone(pixel_values, output_hidden_states=True)
+        layers = [outputs.hidden_states[i] for i in [3, 6, 9, 12]]
+        
+        def process(l, h, w):
+            l = l[:, 1:, :].transpose(1, 2)
+            return l.reshape(l.shape[0], l.shape[1], h//14, w//14)
+            
+        b, _, h, w = pixel_values.shape
+        l1, l2, l3, l4 = [process(layers[i], h, w) for i in range(4)]
+        
+        # Sequential Fusion
+        layer_1_rn = self.scratch.layer1_rn(self.act_postprocess1(l1))
+        layer_2_rn = self.scratch.layer2_rn(self.act_postprocess2(l2))
+        layer_3_rn = self.scratch.layer3_rn(self.act_postprocess3(l3))
+        layer_4_rn = self.scratch.layer4_rn(self.act_postprocess4(l4))
+
+        path_4 = self.refinenet4(layer_4_rn)
+        path_3 = self.refinenet3(path_4, layer_3_rn)
+        path_2 = self.refinenet2(path_3, layer_2_rn)
+        path_1 = self.refinenet1(path_2, layer_1_rn)
+
+        return self.output_conv(path_1).squeeze(1)
+
+AnyThermalDepthModel.register_for_auto_class("AutoModel")
+
 AnyThermalVPRModel.register_for_auto_class("AutoModel")
 # Register for AutoModel discovery
 AnyThermalSegmentationModel.register_for_auto_class("AutoModel")