Upload model

config.json

@@ -1,10 +1,10 @@
 {
   "architectures": [
-    "MambaVisionModelForImageClassification"
+    "MambaVisionModel"
   ],
   "auto_map": {
     "AutoConfig": "configuration_mambavision.MambaVisionConfig",
-    "AutoModelForImageClassification": "modeling_mambavision.MambaVisionModelForImageClassification"
+    "AutoModel": "modeling_mambavision.MambaVisionModel"
   },
   "depths": [
     1,

configuration_mambavision.py

@@ -1,6 +1,4 @@
 from transformers import PretrainedConfig
-from typing import List
-
 
 class MambaVisionConfig(PretrainedConfig):
     model_type = "mambavision"

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9970e390c1e0014ceaee412e485de6016e8db68aef6823dcc2801e838b2be114
+oid sha256:4f6987be0a2ca2222f386eb750d028a05203b047d3c8dfb664c27e2295d02fc0
 size 127219000

modeling_mambavision.py

@@ -28,7 +28,7 @@ from einops import rearrange, repeat
 
 from transformers import PreTrainedModel
 
-from .configuration_mambavision import MambaVisionConfig
+from configuration_mambavision import MambaVisionConfig
 
 
 def _cfg(url='', **kwargs):
@@ -602,8 +602,8 @@ class MambaVisionLayer(nn.Module):
             if pad_r > 0 or pad_b > 0:
                 x = x[:, :, :H, :W].contiguous()
         if self.downsample is None:
-            return x
-        return self.downsample(x)
+            return x, x
+        return self.downsample(x), x
 
 
 class MambaVision(nn.Module):
@@ -697,15 +697,17 @@ class MambaVision(nn.Module):
 
     def forward_features(self, x):
         x = self.patch_embed(x)
+        outs = []
         for level in self.levels:
-            x = level(x)
+            x, xo = level(x)
+            outs.append(xo)
         x = self.norm(x)
         x = self.avgpool(x)
         x = torch.flatten(x, 1)
-        return x
+        return x, outs
 
     def forward(self, x):
-        x = self.forward_features(x)
+        x, outs = self.forward_features(x)
         x = self.head(x)
         return x