Move encoders weights to HF-Hub

segmentation_models_pytorch/encoders/__init__.py

@@ -1,8 +1,12 @@
+import json
 import timm
 import copy
 import warnings
 import functools
-import torch.utils.model_zoo as model_zoo
+from torch.utils.model_zoo import load_url
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
 
 from .resnet import resnet_encoders
 from .dpn import dpn_encoders
@@ -22,6 +26,7 @@
 from .timm_universal import TimmUniversalEncoder
 
 from ._preprocessing import preprocess_input
+from ._legacy_pretrained_settings import pretrained_settings
 
 __all__ = [
     "encoders",
@@ -101,15 +106,40 @@
     encoder = EncoderClass(**params)
 
     if weights is not None:
-        try:
-            settings = encoders[name]["pretrained_settings"][weights]
-        except KeyError:
+        if weights not in encoders[name]["pretrained_settings"]:
+            available_weights = list(encoders[name]["pretrained_settings"].keys())
             raise KeyError(
-                "Wrong pretrained weights `{}` for encoder `{}`. Available options are: {}".format(
-                    weights, name, list(encoders[name]["pretrained_settings"].keys())
-                )
+                f"Wrong pretrained weights `{weights}` for encoder `{name}`. "
+                f"Available options are: {available_weights}"
             )
-        encoder.load_state_dict(model_zoo.load_url(settings["url"]))
+
+        settings = encoders[name]["pretrained_settings"][weights]
+        repo_id = settings["repo_id"]
+        revision = settings["revision"]
+
+        # First, try to load from  HF-Hub, but as far as I know not all countries have
+        # access to the Hub (e.g. China), so we try to load from the original url if
+        # the first attempt fails.
+        try:
+            hf_hub_download(repo_id, filename="config.json", revision=revision)
+            model_path = hf_hub_download(
+                repo_id, filename="model.safetensors", revision=revision
+            )
+            state_dict = load_file(model_path, device="cpu")
+        except Exception as e:
+            if name in pretrained_settings and weights in pretrained_settings[name]:
+                message = (
+                    f"Error loading {name} `{weights}` weights from Hugging Face Hub, "
+                    "trying loading from original url..."
+                )
+                warnings.warn(message, UserWarning)
+                url = pretrained_settings[name][weights]["url"]
+                state_dict = load_url(url, map_location="cpu")
+            else:
+                raise e
+
+        # Load model weights
+        encoder.load_state_dict(state_dict)
 
     encoder.set_in_channels(in_channels, pretrained=weights is not None)
     if output_stride != 32:
@@ -136,7 +166,25 @@
             raise ValueError(
                 "Available pretrained options {}".format(all_settings.keys())
             )
-        settings = all_settings[pretrained]
+
+        repo_id = all_settings[pretrained]["repo_id"]
+        revision = all_settings[pretrained]["revision"]
+
+        # Load config and model
+        try:
+            config_path = hf_hub_download(
+                repo_id, filename="config.json", revision=revision
+            )
+            with open(config_path, "r") as f:
+                settings = json.load(f)
+        except Exception as e:
+            if (
+                encoder_name in pretrained_settings
+                and pretrained in pretrained_settings[encoder_name]
+            ):
+                settings = pretrained_settings[encoder_name][pretrained]
+            else:
+                raise e
 
     formatted_settings = {}
     formatted_settings["input_space"] = settings.get("input_space", "RGB")