Safely load weights from specified URLs

dinov2/hub/backbones.py

@@ -8,9 +8,7 @@
 from typing import Optional, Union
 from urllib.parse import urlparse
 
-import torch
-
-from .utils import _DINOV2_BASE_URL, _make_dinov2_model_name
+from .utils import _DINOV2_BASE_URL, _make_dinov2_model_name, _safe_load_state_dict_from_url
 
 
 class Weights(Enum):
@@ -73,7 +71,7 @@ def _make_dinov2_model(
             url = _DINOV2_BASE_URL + f"/{model_base_name}/{model_full_name}_pretrain.pth"
         else:
             url = convert_path_or_url_to_url(weights)
-        state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu", check_hash=check_hash)
+        state_dict = _safe_load_state_dict_from_url(url, map_location="cpu", check_hash=check_hash)
         model.load_state_dict(state_dict, strict=True)
 
     return model

dinov2/hub/cell_dino/backbones.py

@@ -8,6 +8,8 @@
 
 import torch
 
+from ..utils import _safe_load_state_dict_from_url
+
 
 class Weights(Enum):
     CELL_DINO = "CELL-DINO"
@@ -57,8 +59,8 @@ def _make_cell_dino_model(
         if pretrained_path is not None:
             state_dict = torch.load(pretrained_path, map_location="cpu")
         else:
-            pretrained_url is not None
-            state_dict = torch.hub.load_state_dict_from_url(pretrained_url, map_location="cpu")
+            assert pretrained_url is not None
+            state_dict = _safe_load_state_dict_from_url(pretrained_url, map_location="cpu")
         model.load_state_dict(state_dict, strict=True)
 
     return model

dinov2/hub/classifiers.py

@@ -10,7 +10,7 @@
 import torch.nn as nn
 
 from .backbones import _make_dinov2_model
-from .utils import _DINOV2_BASE_URL, _make_dinov2_model_name
+from .utils import _DINOV2_BASE_URL, _make_dinov2_model_name, _safe_load_state_dict_from_url
 
 
 class Weights(Enum):
@@ -43,7 +43,7 @@ def _make_dinov2_linear_classification_head(
         model_full_name = _make_dinov2_model_name(arch_name, patch_size, num_register_tokens)
         layers_str = str(layers) if layers == 4 else ""
         url = _DINOV2_BASE_URL + f"/{model_base_name}/{model_full_name}_linear{layers_str}_head.pth"
-        state_dict = torch.hub.load_state_dict_from_url(url, map_location="cpu")
+        state_dict = _safe_load_state_dict_from_url(url, map_location="cpu")
         linear_head.load_state_dict(state_dict, strict=True)
 
     return linear_head

dinov2/hub/depthers.py

@@ -7,11 +7,9 @@
 from functools import partial
 from typing import Optional, Tuple, Union
 
-import torch
-
 from .backbones import _make_dinov2_model
 from .depth import BNHead, DepthEncoderDecoder, DPTHead
-from .utils import _DINOV2_BASE_URL, _make_dinov2_model_name, CenterPadding
+from .utils import _DINOV2_BASE_URL, _make_dinov2_model_name, _safe_load_state_dict_from_url, CenterPadding
 
 
 class Weights(Enum):
@@ -132,7 +130,7 @@ def _make_dinov2_linear_depther(
         layers_str = str(layers) if layers == 4 else ""
         weights_str = weights.value.lower()
         url = _DINOV2_BASE_URL + f"/{model_name}/{model_name}_{weights_str}_linear{layers_str}_head.pth"
-        checkpoint = torch.hub.load_state_dict_from_url(url, map_location="cpu")
+        checkpoint = _safe_load_state_dict_from_url(url, map_location="cpu")
         if "state_dict" in checkpoint:
             state_dict = checkpoint["state_dict"]
         model.load_state_dict(state_dict, strict=False)
@@ -220,7 +218,7 @@ def _make_dinov2_dpt_depther(
     if pretrained:
         weights_str = weights.value.lower()
         url = _DINOV2_BASE_URL + f"/{model_name}/{model_name}_{weights_str}_dpt_head.pth"
-        checkpoint = torch.hub.load_state_dict_from_url(url, map_location="cpu")
+        checkpoint = _safe_load_state_dict_from_url(url, map_location="cpu")
         if "state_dict" in checkpoint:
             state_dict = checkpoint["state_dict"]
         model.load_state_dict(state_dict, strict=False)

dinov2/hub/dinotxt.py

@@ -3,11 +3,10 @@
 # This source code is licensed under the Apache License, Version 2.0
 # found in the LICENSE file in the root directory of this source tree.
 
-import torch
 import math
 
 from .backbones import dinov2_vitl14_reg
-from .utils import _DINOV2_BASE_URL
+from .utils import _DINOV2_BASE_URL, _safe_load_state_dict_from_url
 
 
 def dinov2_vitl14_reg4_dinotxt_tet1280d20h24l():
@@ -54,14 +53,10 @@ def dinov2_vitl14_reg4_dinotxt_tet1280d20h24l():
     model.visual_model.backbone = vision_backbone
     model.eval()
 
-    visual_model_head_state_dict = torch.hub.load_state_dict_from_url(
-        _DINOV2_BASE_URL + "/dinov2_vitl14/dinov2_vitl14_reg4_dinotxt_tet1280d20h24l_vision_head.pth",
-        map_location="cpu",
-    )
-    text_model_state_dict = torch.hub.load_state_dict_from_url(
-        _DINOV2_BASE_URL + "/dinov2_vitl14/dinov2_vitl14_reg4_dinotxt_tet1280d20h24l_text_encoder.pth",
-        map_location="cpu",
-    )
+    url = _DINOV2_BASE_URL + "/dinov2_vitl14/dinov2_vitl14_reg4_dinotxt_tet1280d20h24l_vision_head.pth"
+    visual_model_head_state_dict = _safe_load_state_dict_from_url(url, map_location="cpu")
+    url = _DINOV2_BASE_URL + "/dinov2_vitl14/dinov2_vitl14_reg4_dinotxt_tet1280d20h24l_text_encoder.pth"
+    text_model_state_dict = _safe_load_state_dict_from_url(url, map_location="cpu")
     model.visual_model.head.load_state_dict(visual_model_head_state_dict, strict=True)
     model.text_model.load_state_dict(text_model_state_dict, strict=True)
     return model

dinov2/hub/utils.py

@@ -14,6 +14,15 @@
 _DINOV2_BASE_URL = "https://dl.fbaipublicfiles.com/dinov2"
 
 
+def _safe_load_state_dict_from_url(url: str, **kwargs):
+    # See https://github.com/pytorch/pytorch/releases/tag/v2.1.0 (Misc / #98479)
+    if torch.__version__ >= (2, 1):
+        local_kwargs = {**kwargs, "weights_only": True}
+    else:
+        local_kwargs = kwargs
+    return torch.hub.load_state_dict_from_url(url, **local_kwargs)
+
+
 def _make_dinov2_model_name(arch_name: str, patch_size: int, num_register_tokens: int = 0) -> str:
     compact_arch_name = arch_name.replace("_", "")[:4]
     registers_suffix = f"_reg{num_register_tokens}" if num_register_tokens else ""

dinov2/utils/utils.py

@@ -19,7 +19,11 @@
 
 def load_pretrained_weights(model, pretrained_weights, checkpoint_key):
     if urlparse(pretrained_weights).scheme:  # If it looks like an URL
-        state_dict = torch.hub.load_state_dict_from_url(pretrained_weights, map_location="cpu")
+        kwargs = {}
+        # See https://github.com/pytorch/pytorch/releases/tag/v2.1.0 (Misc / #98479)
+        if torch.__version__ >= (2, 1):
+            kwargs["weights_only"] = True
+        state_dict = torch.hub.load_state_dict_from_url(pretrained_weights, map_location="cpu", **kwargs)
     else:
         state_dict = torch.load(pretrained_weights, map_location="cpu")
     if checkpoint_key is not None and checkpoint_key in state_dict: