diff --git a/README.md b/README.md index 588f3ae..9f0c2b5 100644 --- a/README.md +++ b/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 +* Virchow-SPECTRA: https://huggingface.co/medarc/Virchow-SPECTRA * GenBio-PathFM: https://huggingface.co/genbio-ai/genbio-pathfm * CONCH: https://huggingface.co/MahmoodLab/CONCH * TITAN/CONCHv1.5: https://huggingface.co/MahmoodLab/TITAN diff --git a/src/thunder/benchmark.py b/src/thunder/benchmark.py index 01f0e58..969a10f 100644 --- a/src/thunder/benchmark.py +++ b/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*, *virchow_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* diff --git a/src/thunder/config/pretrained_model/virchow_spectra.yaml b/src/thunder/config/pretrained_model/virchow_spectra.yaml new file mode 100644 index 0000000..3a5f22b --- /dev/null +++ b/src/thunder/config/pretrained_model/virchow_spectra.yaml @@ -0,0 +1,12 @@ +model_name: virchow_spectra +type: timm +vlm: false +emb_dim: 2560 +emb_dim_seg: 1280 +ckpt_path: ${oc.env:THUNDER_BASE_DATA_FOLDER}/pretrained_ckpts/virchowspectra/model.safetensors +hf_tag: hf-hub:medarc/Virchow-SPECTRA + +timm_kwargs: + mlp_layer: SwiGLUPacked + act_layer: SiLU + global_pool: "" diff --git a/src/thunder/models/download.py b/src/thunder/models/download.py index db7e381..d5c4261 100644 --- a/src/thunder/models/download.py +++ b/src/thunder/models/download.py @@ -18,6 +18,10 @@ "MahmoodLab/UNI2-h", "pytorch_model.bin", ), # UNI2-h (https://huggingface.co/MahmoodLab/UNI2-h) + "virchow_spectra": ( + "medarc/Virchow-SPECTRA", + "model.safetensors", + ), # Virchow-SPECTRA (https://huggingface.co/medarc/Virchow-SPECTRA) "virchow": ( "paige-ai/Virchow", "pytorch_model.bin", @@ -156,6 +160,7 @@ def download_models(models: Union[List[str], str]) -> None: * uni * uni2h * virchow + * virchow_spectra * virchow2 * hoptimus0 * h0mini @@ -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", "virchowspectra"]: extra_files = ["config.json"] elif local_dir_tag == "titan": extra_files = [ diff --git a/src/thunder/models/pretrained_models.py b/src/thunder/models/pretrained_models.py index f39af21..ec377a4 100644 --- a/src/thunder/models/pretrained_models.py +++ b/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 @@ -327,12 +329,12 @@ def extract_embedding(src, pretrained_model, task_type="linear_probing"): return emb - elif model_cfg.model_name in ["h0mini", "virchow", "virchow2"]: + elif model_cfg.model_name in ["h0mini", "virchow", "virchow2", "virchow_spectra"]: def extract_embedding(src, pretrained_model, task_type): out = pretrained_model(src) class_token = out[:, 0] - if model_cfg.model_name == "virchow": + if model_cfg.model_name in ["virchow", "virchow_spectra"]: patch_tokens = out[:, 1:] elif model_cfg.model_name == "virchow2": patch_tokens = out[ @@ -394,6 +396,7 @@ def get_model_from_name(model_name: str, device: str): * uni * uni2h * virchow + * virchow_spectra * virchow2 * hoptimus0 * h0mini @@ -470,16 +473,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( diff --git a/src/thunder/utils/constants.py b/src/thunder/utils/constants.py index 5a71411..b535a2c 100644 --- a/src/thunder/utils/constants.py +++ b/src/thunder/utils/constants.py @@ -21,6 +21,7 @@ class ModelConstants(Enum): "uni", "uni2h", "virchow", + "virchow_spectra", "virchow2", "provgigapath", # Histo-image-language models