Update model/run_inference.py

model/run_inference.py

@@ -50,17 +50,17 @@ CAM_ORDER = [
 DUMMY_4x4 = np.eye(4, dtype=np.float32)
 
 
-def _download_ckpt_from_hf(repo_id: str, filename: str, cache_dir: str = "./.hf_cache") -> str:
-    """
-    Download a checkpoint from Hugging Face Hub. Uses HF_TOKEN env if set (for private repos).
-    """
+def _download_ckpt_from_hf(repo_id: str, filename: str, cache_dir: str = None) -> str:
+    # default to a writable cache location in Spaces/containers
+    cache_dir = cache_dir or os.environ.get("HF_HUB_CACHE") or "/tmp/hf_cache"
+    os.makedirs(cache_dir, exist_ok=True)
+
     ckpt_path = hf_hub_download(
         repo_id=repo_id,
         filename=filename,
-        revision="main",           # change to a tag/commit if you want pinning
-        cache_dir=cache_dir,       # keeps the file across runs
-        local_dir=None,            # rely on cache_dir
-        token=os.environ.get("HF_TOKEN", None),
+        revision="main",
+        cache_dir=cache_dir,
+        token=os.environ.get("HF_TOKEN"),
     )
     if not os.path.isfile(ckpt_path):
         raise FileNotFoundError(f"Failed to download checkpoint: {repo_id}/{filename}")