Updating model

model_loader.py

@@ -1,3 +1,4 @@
+import os
 import torch
 import torch.nn as nn
 from huggingface_hub import hf_hub_download
@@ -24,7 +25,22 @@ def build_classifier(num_classes: int,
     device = torch.device(device)
 
     # 1) Download pretrained MAE weights from the Hub
-    ckpt_path = hf_hub_download(repo_id=base_repo, filename=base_filename)
+    #ckpt_path = hf_hub_download(repo_id=base_repo, filename=base_filename)
+    # Read token from env (if set)
+    hf_token = os.getenv("HF_TOKEN")
+
+    try:
+        ckpt_path = hf_hub_download(
+            repo_id=base_repo,
+            filename=base_filename,
+            token=hf_token,               # Works for private if token exists
+            cache_dir="/tmp/hf_cache"      # Spaces-friendly cache
+        )
+    except Exception as e:
+        raise RuntimeError(f"Failed to download model from {base_repo}: {e}")
+
+    # Load model weights
+    print(f"Loading RETFound MAE weights from {ckpt_path}...")
 
     # 2) Build backbone
     model = RETFound_mae(global_pool=global_pool, drop_path_rate=drop_path_rate)