Upload eigen_moe.py with huggingface_hub

eigen_moe.py

@@ -390,13 +390,13 @@ class HFEigenMoE(nn.Module, PyTorchModelHubMixin):
         vit_model_name: str,
     ) -> str:
         base = Path(model_dir)
-        candidates = []
         if checkpoint_filename:
-            candidates.append(checkpoint_filename)
-        default_name = default_hub_checkpoint_filename(vit_model_name)
-        if default_name:
-            candidates.append(default_name)
-        candidates.extend(["model.safetensors", "pytorch_model.bin"])
+            candidates = [checkpoint_filename]
+        else:
+            candidates = ["model.safetensors", "pytorch_model.bin"]
+            default_name = default_hub_checkpoint_filename(vit_model_name)
+            if default_name:
+                candidates.append(default_name)
 
         for filename in candidates:
             path = base / filename
@@ -429,13 +429,13 @@ class HFEigenMoE(nn.Module, PyTorchModelHubMixin):
         if hf_hub_download is None:
             raise ImportError("huggingface_hub is required to download checkpoints from the Hub.")
 
-        candidates = []
         if checkpoint_filename:
-            candidates.append(checkpoint_filename)
-        default_name = default_hub_checkpoint_filename(vit_model_name)
-        if default_name:
-            candidates.append(default_name)
-        candidates.extend(["model.safetensors", "pytorch_model.bin"])
+            candidates = [checkpoint_filename]
+        else:
+            candidates = ["model.safetensors", "pytorch_model.bin"]
+            default_name = default_hub_checkpoint_filename(vit_model_name)
+            if default_name:
+                candidates.append(default_name)
 
         seen = set()
         unique_candidates = []