Add OpenMidnight-1.1-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
+* OpenMidnight-1.1-SPECTRA: https://huggingface.co/medarc/OpenMidnight-1.1-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*, *openmidnight11_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/openmidnight11_spectra.yaml

@@ -0,0 +1,8 @@
+model_name: openmidnight11_spectra
+type: timm
+vlm: false
+emb_dim: 1536
+ckpt_path: ${oc.env:THUNDER_BASE_DATA_FOLDER}/pretrained_ckpts/openmidnight1.1spectra/model.safetensors
+hf_tag: hf-hub:medarc/OpenMidnight-1.1-SPECTRA
+
+timm_kwargs: {}

src/thunder/models/download.py

@@ -90,6 +90,10 @@
         "kaiko-ai/midnight",
         "model.safetensors",
     ),  # Midnight-12k (https://huggingface.co/kaiko-ai/midnight)
+    "openmidnight11_spectra": (
+        "medarc/OpenMidnight-1.1-SPECTRA",
+        "model.safetensors",
+    ),  # OpenMidnight-1.1-SPECTRA (https://huggingface.co/medarc/OpenMidnight-1.1-SPECTRA)
     "openmidnight": (
         "SophontAI/OpenMidnight",
         "teacher_checkpoint_load.pt",
@@ -170,6 +174,7 @@ def download_models(models: Union[List[str], str]) -> None:
         * hiboul
         * midnight
         * openmidnight
+        * openmidnight11_spectra
         * kaiko_vits8
         * kaiko_vits16
         * kaiko_vitb8
@@ -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", "openmidnight1.1spectra"]:
                 extra_files = ["config.json"]
             elif local_dir_tag == "titan":
                 extra_files = [

src/thunder/models/pretrained_models.py

@@ -1,9 +1,11 @@
+import json
 from abc import ABC, abstractmethod
 from pathlib import Path
 
 import timm
 import torch
 from omegaconf import DictConfig, OmegaConf
+from safetensors.torch import load_file
 from timm.data import resolve_data_config
 from timm.data.transforms_factory import create_transform
 from transformers import AutoImageProcessor, AutoModel, AutoTokenizer
@@ -375,7 +377,11 @@ def extract_embedding(src, pretrained_model, task_type="linear_probing"):
 
                     if model_cfg.model_name == "uni2h":
                         nb_reg = 8
-                    elif model_cfg.model_name in ["hoptimus0", "hoptimus1"]:
+                    elif model_cfg.model_name in [
+                        "hoptimus0",
+                        "hoptimus1",
+                        "openmidnight11_spectra",
+                    ]:
                         nb_reg = 4
                     else:
                         nb_reg = 0
@@ -408,6 +414,7 @@ def get_model_from_name(model_name: str, device: str):
         * hiboul
         * midnight
         * openmidnight
+        * openmidnight11_spectra
         * kaiko_vits8
         * kaiko_vits16
         * kaiko_vitb8
@@ -470,16 +477,28 @@ def get_from_timm(hf_tag: str, timm_kwargs: dict, ckpt_path: str, device: str):
 
     :param hf_tag: Huggingface model tag.
     :param timm_kargs: dictionary of timm arguments.
-    :ckpt_path: path to the stored checkpoint.
+    :param ckpt_path: path to the stored checkpoint. Timm safetensors require
+        an adjacent config.json.
     :param device: device to use (cpu, cuda).
     """
 
     # Model
-    model = timm.create_model(hf_tag, pretrained=False, **timm_kwargs)
-    model.load_state_dict(
-        torch.load(ckpt_path, weights_only=True, map_location=torch.device(device)),
-        strict=True,
-    )
+    if Path(ckpt_path).suffix == ".safetensors":
+        # Read the downloaded timm configuration so loading also works offline.
+        config = json.loads(Path(ckpt_path).with_name("config.json").read_text())
+        model = timm.create_model(
+            config["architecture"],
+            pretrained=False,
+            pretrained_cfg=config["pretrained_cfg"],
+            **(config["model_args"] | timm_kwargs),
+        )
+        model.load_state_dict(load_file(ckpt_path, device=device), strict=True)
+    else:
+        model = timm.create_model(hf_tag, pretrained=False, **timm_kwargs)
+        model.load_state_dict(
+            torch.load(ckpt_path, weights_only=True, map_location=torch.device(device)),
+            strict=True,
+        )
 
     # Transform
     transform = create_transform(

src/thunder/utils/constants.py

@@ -16,6 +16,7 @@ class ModelConstants(Enum):
         "kaiko_vitb16",
         "midnight",
         "openmidnight",
+        "openmidnight11_spectra",
         "phikon",
         "phikon2",
         "uni",