Fix Space model downloads and torch deps

Add a Hugging Face Spaces pre-requirements file to force a newer pip before installing app dependencies, and switch PyTorch model loading to download SuperAnimal weights into a writable local directory instead of the read-only site-packages model cache. The requirements update also pins Spaces to CPU-only PyTorch wheels to avoid pulling unnecessary CUDA packages.

pre-requirements.txt

@@ -0,0 +1 @@
+pip>=26.2

pytorch_utils.py

@@ -1,10 +1,12 @@
 import threading
+from pathlib import Path
 
 import numpy as np
 import PIL
 from deeplabcut.pose_estimation_pytorch.apis.utils import get_inference_runners
 from deeplabcut.pose_estimation_pytorch.config.pose import PoseConfig
-from deeplabcut.pose_estimation_pytorch.modelzoo.utils import get_super_animal_snapshot_path
+from deeplabcut.pose_estimation_pytorch.modelzoo.utils import MODEL_FILENAME_MAPPING
+from dlclibrary import download_huggingface_model
 
 # SuperAnimal (pose model, detector) used by the PyTorch backend
 PYTORCH_MODELS = {
@@ -14,17 +16,28 @@ PYTORCH_MODELS = {
 
 MAX_INDIVIDUALS = 10
 MAX_IMAGE_SIZE = 1280  # longest side fed to the models (and drawn on)
+# next to the TF models, not deeplabcut/modelzoo/checkpoints: site-packages is read-only for the Space's non-root user
+WEIGHTS_DIR = Path(__file__).parent / "DLC_models" / "pytorch"
 
 _runners = {}
 _build_lock = threading.Lock()
 
 
 ##########################################
-def load_superanimal(superanimal, device="auto"):
-    """Build (once) the detector and pose runners for a SuperAnimal model.
+def snapshot_path(superanimal, model_name):
+    """Path to a SuperAnimal snapshot in WEIGHTS_DIR, downloaded on first use (as deeplabcut does)."""
+    name = f"{superanimal}_{model_name}"
+    path = WEIGHTS_DIR / f"{name}.pt"
+    if not path.exists():
+        source = MODEL_FILENAME_MAPPING.get(name, path.name)
+        rename = None if source == path.name else {source: path.name}
+        download_huggingface_model(name, target_dir=str(WEIGHTS_DIR), rename_mapping=rename)
+    return path
 
-    Weights are downloaded on first use to deeplabcut/modelzoo/checkpoints.
-    """
+
+##########################################
+def load_superanimal(superanimal, device="auto"):
+    """Build (once) the detector and pose runners for a SuperAnimal model; weights are downloaded on first use."""
     with _build_lock:
         if superanimal not in _runners:
             pose_model, detector = PYTORCH_MODELS[superanimal]
@@ -39,8 +52,8 @@ def load_superanimal(superanimal, device="auto"):
             cfg["detector"]["model"]["box_score_thresh"] = 0.05
             pose_runner, detector_runner = get_inference_runners(
                 cfg,
-                snapshot_path=get_super_animal_snapshot_path(superanimal, pose_model),
-                detector_path=get_super_animal_snapshot_path(superanimal, detector),
+                snapshot_path=snapshot_path(superanimal, pose_model),
+                detector_path=snapshot_path(superanimal, detector),
                 max_individuals=MAX_INDIVIDUALS,
                 inference_cfg={"multithreading": {"enabled": False}},
             )

requirements.txt

@@ -1,4 +1,6 @@
 # Hugging Face Spaces install from this file; keep in sync with pyproject.toml dependencies.
+# CPU-only PyTorch wheels: the Space has no GPU, and the default Linux wheels pull ~2-3 GB of CUDA libraries
+--extra-index-url https://download.pytorch.org/whl/cpu
 gradio==6.29.0
 gitpython>=3.1.30
 seaborn