diff --git a/src/eva/vision/models/networks/backbones/pathology/mahmood.py b/src/eva/vision/models/networks/backbones/pathology/mahmood.py index 5aa55344..2bfa8717 100644 --- a/src/eva/vision/models/networks/backbones/pathology/mahmood.py +++ b/src/eva/vision/models/networks/backbones/pathology/mahmood.py @@ -1,13 +1,9 @@ """Pathology FMs from MahmoodLab.""" -import os -from pathlib import Path from typing import Tuple -import huggingface_hub import timm import torch -from loguru import logger from torch import nn from eva.vision.models import wrappers @@ -20,7 +16,6 @@ def mahmood_uni( dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None, hf_token: str | None = None, - download_dir: str = os.path.join(str(Path.home()), ".cache/eva"), ) -> nn.Module: """Initializes UNI model from MahmoodLab. @@ -29,31 +24,20 @@ def mahmood_uni( the grid size (interpolate abs and/or ROPE pos) in the forward pass. out_indices: Whether and which multi-level patch embeddings to return. hf_token: HuggingFace token to download the model. - download_dir: Directory to download the model checkpoint. Returns: The model instance. """ - checkpoint_path = os.path.join(download_dir, "uni.bin") - if not os.path.exists(checkpoint_path): - logger.info(f"Downloading the model checkpoint to {download_dir} ...") - os.makedirs(download_dir, exist_ok=True) - _utils.huggingface_login(hf_token) - huggingface_hub.hf_hub_download( - "MahmoodLab/UNI", - filename="uni.bin", - local_dir=download_dir, - force_download=True, - ) + _utils.huggingface_login(hf_token) return wrappers.TimmModel( - model_name="vit_large_patch16_224", + model_name="hf-hub:MahmoodLab/uni", + pretrained=True, out_indices=out_indices, model_kwargs={ "init_values": 1e-5, "dynamic_img_size": dynamic_img_size, }, - checkpoint_path=checkpoint_path, )