added CAPI MIM model

README.md

@@ -78,7 +78,8 @@ Neural networks come from different sources. With `thingsvision`, you can extrac
   - `dino-rn50`, `dino-xcit-{small/medium}-{12/24}-p{8/16}`
   - `dino-vit-{tiny/small/base}-p{8/16}`
   - `dinov2-vit-{small/base/large/giant}-p14`
-  - `mae-vit-{base/large}-p16`, `mae-vit-huge-p14`<br>
+  - `mae-vit-{base/large}-p16`, `mae-vit-huge-p14`
+  - `capi-vitl14-{p205/lvd/in22k/in1k}` (trained on different datasets)<br>
 - [OpenCLIP](https://github.com/mlfoundations/open_clip) models (CLIP trained on LAION-{400M/2B/5B})
 - [CLIP](https://github.com/openai/CLIP) models (CLIP trained on WiT)
 - a few custom models (Alexnet, VGG-16, Resnet50, and Inception_v3) trained on [Ecoset](https://www.pnas.org/doi/10.1073/pnas.2011417118)<br>

requirements.txt

@@ -12,7 +12,7 @@ scikit-learn
 scipy
 tensorflow<2.16
 timm
-torch>=2.0.0
+torch>=2.4.0
 torchvision==0.15.2
 torchtyping
 tqdm
@@ -24,3 +24,7 @@ keras-cv-attention-models>=1.3.5
 vit-keras==0.1.2
 git+https://github.com/serre-lab/Harmonization.git
 dreamsim==0.1.3
+jaxtyping
+omegaconf
+einops
+rich

tests/helper.py

@@ -219,6 +219,13 @@
         "source": "ssl",
         "kwargs": {"token_extraction": "avg_pool"},
     },
+    "capi-vitl14-in1k": {
+        "model_name": "capi-vitl14-in1k",
+        "modules": ["norm", "fc_norm"],
+        "pretrained": True,
+        "source": "ssl",
+        "kwargs": {"token_extraction": "cls_token"},
+    },
     # Additional models
     "Harmonization_visual_ResNet50": {
         "model_name": "Harmonization",

thingsvision/core/extraction/extractors.py

@@ -383,6 +383,26 @@ class SSLExtractor(PyTorchExtractor):
             "type": "hub",
             "checkpoint_url": "https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_huge.pth",
         },
+        "capi-vitl14-p205": {
+            "repository": "facebookresearch/capi:main",
+            "arch": "capi_vitl14_p205",
+            "type": "hub",
+        },
+        "capi-vitl14-lvd": {
+            "repository": "facebookresearch/capi:main",
+            "arch": "capi_vitl14_lvd",
+            "type": "hub",
+        },
+        "capi-vitl14-in22k": {
+            "repository": "facebookresearch/capi:main",
+            "arch": "capi_vitl14_in22k",
+            "type": "hub",
+        },
+        "capi-vitl14-in1k": {
+            "repository": "facebookresearch/capi:main",
+            "arch": "capi_vitl14_in1k",
+            "type": "hub",
+        },
     }
 
     def __init__(