Add Midnight-12k-SPECTRA encoder

README.md

@@ -65,6 +65,7 @@ An API and command line interface (CLI) are provided to allow users to download
 * Virchow2: https://huggingface.co/paige-ai/Virchow2
 * H-optimus-0: https://huggingface.co/bioptimus/H-optimus-0
 * H-optimus-1: https://huggingface.co/bioptimus/H-optimus-1
+* Midnight-12k-SPECTRA: https://huggingface.co/medarc/Midnight-12k-SPECTRA
 * GenBio-PathFM: https://huggingface.co/genbio-ai/genbio-pathfm
 * CONCH: https://huggingface.co/MahmoodLab/CONCH
 * TITAN/CONCHv1.5: https://huggingface.co/MahmoodLab/TITAN

src/thunder/benchmark.py

@@ -21,7 +21,7 @@ def benchmark(
 
     where options are:
         - dataset: *bach*, *bracs*, *break_his*, *ccrcc*, *crc*, *esca*, *mhist*, *ocelot*, *pannuke*, *patch_camelyon*, *segpath_epithelial*, *segpath_lymphocytes*, *tcga_crc_msi*, *tcga_tils*, *tcga_uniform*, *wilds*
-        - model: *hiboub*, *hiboul*, *hoptimus0*, *hoptimus1*, *genbio-pathfm*, *midnight*, *phikon*, *phikon2*, *uni*, *uni2h*, *virchow*, *virchow2*, *conch*, *titan*, *keep*, *musk*, *plip*, *quiltnetb32*, *dinov2base*, *dinov2large*, *vitbasepatch16224in21k*, *vitlargepatch16224in21k*, *clipvitbasepatch32*, *clipvitlargepatch14*
+        - model: *hiboub*, *hiboul*, *hoptimus0*, *hoptimus1*, *genbio-pathfm*, *midnight12k_spectra*, *midnight*, *phikon*, *phikon2*, *uni*, *uni2h*, *virchow*, *virchow2*, *conch*, *titan*, *keep*, *musk*, *plip*, *quiltnetb32*, *dinov2base*, *dinov2large*, *vitbasepatch16224in21k*, *vitlargepatch16224in21k*, *clipvitbasepatch32*, *clipvitlargepatch14*
         - task: *adversarial_attack*, *alignment_scoring*, *image_retrieval*, *knn*, *linear_probing*, *pre_computing_embeddings*, *segmentation*, *simple_shot*, *transformation_invariance*, *zero_shot_vlm*
         - loading_mode: *online_loading*, *image_pre_loading*, *embedding_pre_loading*
 

src/thunder/config/pretrained_model/midnight12k_spectra.yaml

@@ -0,0 +1,6 @@
+model_name: midnight12k_spectra
+type: safetensors
+vlm: false
+emb_dim: 3072
+emb_dim_seg: 1536
+ckpt_path: ${oc.env:THUNDER_BASE_DATA_FOLDER}/pretrained_ckpts/midnight12kspectra/

src/thunder/models/download.py

@@ -86,6 +86,10 @@
         "histai/hibou-L",
         "model.safetensors",
     ),  # Hibou-L (https://huggingface.co/histai/hibou-L)
+    "midnight12k_spectra": (
+        "medarc/Midnight-12k-SPECTRA",
+        "model.safetensors",
+    ),  # Midnight-12k-SPECTRA (https://huggingface.co/medarc/Midnight-12k-SPECTRA)
     "midnight": (
         "kaiko-ai/midnight",
         "model.safetensors",
@@ -169,6 +173,7 @@ def download_models(models: Union[List[str], str]) -> None:
         * hiboub
         * hiboul
         * midnight
+        * midnight12k_spectra
         * openmidnight
         * kaiko_vits8
         * kaiko_vits16
@@ -262,7 +267,7 @@ def download_model(model: str) -> None:
                 ]
                 if local_dir_tag != "plip":
                     extra_files.extend(["merges.txt", "vocab.json"])
-            elif local_dir_tag == "midnight":
+            elif local_dir_tag in ["midnight", "midnight12kspectra"]:
                 extra_files = ["config.json"]
             elif local_dir_tag == "titan":
                 extra_files = [

src/thunder/models/pretrained_models.py

@@ -124,7 +124,7 @@ def get_model(model_cfg: dict, device: str):
             model, transform, tokenizer = get_keep(model_cfg.ckpt_path)
         elif model_cfg.model_name == "titan":
             model, transform = get_titan(model_cfg.ckpt_path)
-        elif model_cfg.model_name == "midnight":
+        elif model_cfg.model_name in ["midnight", "midnight12k_spectra"]:
             model, transform = get_midnight(model_cfg.ckpt_path)
         elif model_cfg.model_name == "genbio-pathfm":
             model, transform = get_genbio_pathfm(model_cfg.ckpt_path)
@@ -346,7 +346,7 @@ def extract_embedding(src, pretrained_model, task_type):
             else:
                 return patch_tokens
 
-    elif model_cfg.model_name == "midnight":
+    elif model_cfg.model_name in ["midnight", "midnight12k_spectra"]:
 
         def extract_embedding(src, pretrained_model, task_type):
             out = pretrained_model(src).last_hidden_state
@@ -407,6 +407,7 @@ def get_model_from_name(model_name: str, device: str):
         * hiboub
         * hiboul
         * midnight
+        * midnight12k_spectra
         * openmidnight
         * kaiko_vits8
         * kaiko_vits16

src/thunder/utils/constants.py

@@ -15,6 +15,7 @@ class ModelConstants(Enum):
         "kaiko_vitb8",
         "kaiko_vitb16",
         "midnight",
+        "midnight12k_spectra",
         "openmidnight",
         "phikon",
         "phikon2",