Auth check for HF repos + better errors

.gitignore

@@ -19,6 +19,7 @@ debug*.py
 config.json
 secrets.json
 *.zip
+training_presets/*.json
 
 # environments
 /.venv*

modules/modelLoader/GenericEmbeddingModelLoader.py

@@ -45,6 +45,7 @@ def load(
             model.model_spec = self._load_default_model_spec(model_type)
 
             if model_names.base_model is not None:
+                base_model_loader._validate_hf_repo_access(model_names.base_model)
                 base_model_loader.load(model, model_type, model_names, weight_dtypes, quantization)
             embedding_loader.load(model, model_names.embedding.model_name, model_names)
 

modules/modelLoader/GenericFineTuneModelLoader.py

@@ -50,6 +50,7 @@ def load(
             self._load_internal_data(model, model_names.base_model)
             model.model_spec = self._load_default_model_spec(model_type)
 
+            base_model_loader._validate_hf_repo_access(model_names.base_model)
             base_model_loader.load(model, model_type, model_names, weight_dtypes, quantization)
             if embedding_loader_class is not None:
                 embedding_loader.load(model, model_names.base_model, model_names)

modules/modelLoader/GenericLoRAModelLoader.py

@@ -48,6 +48,7 @@ def load(
             model.model_spec = self._load_default_model_spec(model_type)
 
             if model_names.base_model is not None:
+                base_model_loader._validate_hf_repo_access(model_names.base_model)
                 base_model_loader.load(model, model_type, model_names, weight_dtypes, quantization)
             lora_model_loader.load(model, model_names)
             if embedding_loader_class is not None:

modules/modelLoader/mixin/HFModelLoaderMixin.py

@@ -19,7 +19,10 @@
 
 import accelerate
 import huggingface_hub
-from huggingface_hub.utils import EntryNotFoundError
+from huggingface_hub import constants as hf_constants
+from huggingface_hub import parse_hf_uri
+from huggingface_hub.errors import HfUriError
+from huggingface_hub.utils import EntryNotFoundError, HfHubHTTPError, RepositoryNotFoundError
 from safetensors.torch import load_file
 
 # huggingface_hub 1.16+ uses httpx, which logs every HTTP request/response at INFO level.
@@ -29,6 +32,54 @@
 class HFModelLoaderMixin(metaclass=ABCMeta):
     def __init__(self):
         super().__init__()
+        self.__validated_hf_repositories: set[str] = set()
+
+    @staticmethod
+    def __normalize_hf_repo_id(repo_id: str) -> str:
+        if not repo_id.startswith(("http://", "https://")):
+            return repo_id
+
+        try:
+            uri = parse_hf_uri(repo_id)
+        except HfUriError:
+            return repo_id
+
+        if uri.type != "model":
+            raise ValueError(f"Expected a Hugging Face model URL, got a {uri.type} URL.")
+
+        return uri.id
+
+    def _validate_hf_repo_access(
+            self,
+            repo_id: str | None,
+    ):
+        if not repo_id \
+                or os.path.exists(repo_id) \
+                or hf_constants.HF_HUB_OFFLINE:
+            return
+
+        repo_id = self.__normalize_hf_repo_id(repo_id)
+        if repo_id in self.__validated_hf_repositories:
+            return
+
+        token = huggingface_hub.get_token()
+        if token is not None:
+            try:
+                huggingface_hub.whoami(token=token, cache=True)
+            except HfHubHTTPError as e:
+                if e.response.status_code == 401:
+                    raise ValueError("Invalid Hugging Face token.") from None
+                raise
+
+        try:
+            huggingface_hub.auth_check(
+                repo_id=repo_id,
+                repo_type="model",
+                token=token,
+            )
+        except RepositoryNotFoundError as e:
+            raise e.with_traceback(None) from None
+        self.__validated_hf_repositories.add(repo_id)
 
     def __load_sub_module(
             self,
@@ -198,6 +249,8 @@ def _load_transformers_sub_module(
             pretrained_model_name_or_path: str,
             subfolder: str = "",
     ):
+        self._validate_hf_repo_access(pretrained_model_name_or_path)
+
         user_agent = {
             "file_type": "model",
             "framework": "pytorch",
@@ -235,6 +288,8 @@ def _load_diffusers_sub_module(
             subfolder: str | None = None,
             quantization: QuantizationConfig | None = None,
     ):
+        self._validate_hf_repo_access(pretrained_model_name_or_path)
+
         user_agent = {
             "file_type": "model",
             "framework": "pytorch",