diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 0000000..c9aeeaf
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1,5 @@
+# Hugging Face Hub stores these via Git LFS / Xet (plain PNG/JPG in git are rejected on push).
+demo_data/*.png filter=lfs diff=lfs merge=lfs -text
+demo_data/*.jpg filter=lfs diff=lfs merge=lfs -text
+demo_data/*.jpeg filter=lfs diff=lfs merge=lfs -text
+images/*.png filter=lfs diff=lfs merge=lfs -text
diff --git a/.gitignore b/.gitignore
index f8b5caa..efea017 100644
--- a/.gitignore
+++ b/.gitignore
@@ -165,6 +165,8 @@ cython_debug/
 
 # Directory
 .gradio/
+demo_data/*.mp4
+*.mp4
 demo_out/
 demo_out*/
 data/PRIMA*/
diff --git a/LICENSE.md b/LICENSE.md
new file mode 100644
index 0000000..ba0a28b
--- /dev/null
+++ b/LICENSE.md
@@ -0,0 +1,13 @@
+Modified MIT.
+
+Copyright 2026 by Mackenzie Mathis, Xiaohang Yu, and contributors.
+
+A fully-paid, non-exclusive, and non-transferable license is hereby granted to you (hereafter "LICENSEE") for academic, non-commercial purposes only (hereafter "LICENSE") to use the "MODEL" weights (hereafter "MODEL"), subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the MODEL.
+
+The MODEL may not be used to deliberately harm any animal.
+
+LICENSEE acknowledges that the MODEL is a research tool. THE MODEL IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE MODEL OR THE USE OR OTHER DEALINGS IN THE MODEL.
+
+If this license is not appropriate for your application, please contact Prof. Mackenzie W. Mathis (mackenzie@post.harvard.edu) and/or the TTO office at EPFL (tto@epfl.ch) for a commercial use license.
diff --git a/README.md b/README.md
index ed0ea58..3787d02 100644
--- a/README.md
+++ b/README.md
@@ -12,20 +12,28 @@ Xiaohang Yu, Ti Wang, Mackenzie Weygandt Mathis
 ---
 
 
-## 🚀 TL;DR
-PRIMA creates a 3D quadruped mesh from a single 2D image. It leverages BioCLIP-based biological priors for robust cross-species shape understanding, then applies test-time adaptation with 2D reprojection and auxiliary keypoint guidance to refine SMAL pose and shape predictions. It further uses this adaptation pipeline to build Quadruped3D, a large-scale pseudo-3D dataset with diverse species and poses, achieving state-of-the-art results on Animal3D, CtrlAni3D, Quadruped2D, and Animal Kingdom datasets.
+## TL;DR
+PRIMA creates a 3D quadruped mesh from a single 2D image. It leverages BioCLIP-based biological priors for robust cross-species shape understanding, then applies test-time adaptation with 2D reprojection and auxiliary keypoint guidance to refine SMAL pose and shape predictions. 
+
+It further can be used to build Quadruped3D, a large-scale pseudo-3D dataset with diverse species and poses. 
+
+PRIMA achieves state-of-the-art results on Animal3D, CtrlAni3D, Quadruped2D, and Animal Kingdom datasets.
 
 ## Installation
 
+PRIMA requires Python 3.10 or newer. A CUDA-enabled PyTorch installation is
+recommended for local inference and training.
+
 ### Install from PyPI
 
-> Recommended: Python 3.10 and a CUDA-enabled PyTorch installation.
+Create a clean environment, install PyTorch for your CUDA version, then install
+the package:
 
 ```bash
 conda create -n prima python=3.10 -y
 conda activate prima
 
-# Install PyTorch matching your CUDA (example: CUDA 11.8)
+# Example for CUDA 11.8. Adjust this command for your CUDA version.
 pip install --index-url https://download.pytorch.org/whl/cu118 \
     "torch==2.2.1" "torchvision==0.17.1" "torchaudio==2.2.1"
 
@@ -35,84 +43,119 @@ python -m pip install --no-build-isolation \
 python -m pip install --no-build-isolation \
       "git+https://github.com/facebookresearch/pytorch3d.git"
 
+# Install PRIMA from PyPI-test (for now)
+pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple prima-animal==0.1.7
+
 # Install PRIMA from PyPI
 pip install prima-animal
 ```
 
 `prima-animal` includes demo runtime dependencies used by `demo.py`, `demo_tta.py`, and `app.py` (including Detectron2 and DeepLabCut).
 
----
+### Install from this repository
 
-## Demo
+Use this path if you want to run the code from a fresh clone.
 
-### Checkpoints and data
+```bash
+git clone https://github.com/AdaptiveMotorControlLab/PRIMA.git
+cd PRIMA
+```
+
+The helper script below creates a fresh virtual environment, installs runtime
+dependencies, pulls Git LFS assets if available, downloads the default demo
+checkpoints/data, and verifies that the demo dependencies can be imported:
 
-Create a `data/` folder under the project root and download the required files into it:
 
 ```bash
-mkdir -p data/smal
+PRIMA_PYTHON=/path/to/python3.10 \
+PRIMA_VENV=prima_env \
+./scripts/clean_install_local.sh
+source prima_env/bin/activate
 ```
 
-1. **SMAL model** -- download from [here](https://drive.google.com/drive/folders/1O1tWYimVMA7hEbnwuPyiDWh90tUGoTPB?usp=drive_link) and place the `.pkl` files under `data/smal/`
-2. **Pretrained backbone** -- download from [here](https://drive.google.com/file/d/1jOJXJVPXnWX7W7vqYVt0joJZr4C8x-Yo/view?usp=drive_link) and place at `data/amr_vitbb.pth`
-3. **Stage-1 checkpoint** -- download from [here](https://drive.google.com/drive/folders/1pwIpYwP3aJ6W2M3-WhEvcFjW38-4j405?usp=drive_link) and place under `data/PRIMAS1/`
-4. **Stage-3 checkpoint** -- download from [here](https://drive.google.com/drive/folders/1DO6idTCORL5G6PLjikRaIjCXmo_-Ut31?usp=drive_link) and place under `data/PRIMAS3/`
+Options:
 
-After downloading, the expected layout is:
+- `--skip-data` skips the large demo data download if `data/` is already populated.
+- `--wipe-data --force-data` removes downloaded demo assets and downloads them again.
+- `--no-editable` installs dependencies without registering the repo as an editable package.
 
+On macOS, install Python 3.10 if needed:
+
+```bash
+brew install python@3.10
+PRIMA_PYTHON=/opt/homebrew/bin/python3.10 \
+PRIMA_VENV=prima_env \
+./scripts/clean_install_local.sh
+source prima_env/bin/activate
 ```
-PRIMA/
-└── data/
-    ├── smal/
-    │   ├── my_smpl_00781_4_all.pkl
-    │   ├── my_smpl_data_00781_4_all.pkl
-    │   └── walking_toy_symmetric_pose_prior_with_cov_35parts.pkl
-    ├── amr_vitbb.pth
-    ├── PRIMAS1/
-    │   ├── .hydra/
-    │   │   └── config.yaml
-    │   └── checkpoints/
-    │       └── s1ckpt.ckpt
-    └── PRIMAS3/
-        ├── .hydra/
-        │   └── config.yaml
-        └── checkpoints/
-            └── s3ckpt.ckpt
+
+If macOS reports `Cannot read image: demo_data/...`, install Git LFS and pull
+the demo images:
+
+```bash
+git lfs install
+git lfs pull --include="demo_data/*"
 ```
 
 ---
 
+## Demo
+
+### Checkpoints and data
+
+The demo scripts auto-download their default Stage 1 PRIMA assets from Hugging
+Face when the checkpoint or matching Hydra config is missing. If you want to
+pre-download all necessary checkpoints and data ahead of time, run:
+
+```bash
+python scripts/setup_demo_data.py --hf-repo-id MLAdaptiveIntelligence/PRIMA
+```
+
+Approximate default prefetch volume from Hugging Face is ~5.5 GB total
+(`s1ckpt_inference.ckpt` ~3 GB + `amr_vitbb.pth` ~2.5 GB + SMAL files).
+Expected time is roughly:
+- 100 Mbps: ~7-10 minutes
+- 300 Mbps: ~2-4 minutes
+- 1 Gbps: ~1 minute
+
+Existing files are reused by default; pass `--force` only if you need to redownload them. If you also need the Stage 3 pretrained model, add `--include-stage3`.
+
+Expected files in that Hugging Face repo root:
+- `my_smpl_00781_4_all.pkl`
+- `my_smpl_data_00781_4_all.pkl`
+- `walking_toy_symmetric_pose_prior_with_cov_35parts.pkl`
+- `amr_vitbb.pth`
+- `config_s1_HYDRA.yaml`
+- `s1ckpt_inference.ckpt`
+
+Optional Stage 3 prefetch expects:
+- `config_s3_HYDRA.yaml`
+- `s3ckpt_inference.ckpt`
+
 ### Demo (without TTA)
 
 Run animal detection + PRIMA 3D pose/shape inference:
 
 ```bash
-python demo.py \
-  --checkpoint data/PRIMAS1/checkpoints/s1ckpt.ckpt \
-  --img_folder demo_data/ \
-  --out_folder demo_out/
+bash demo.sh
 ```
 
-Outputs are written to `demo_out/`.
+Outputs are written to `demo_out/`. Edit `demo.sh` if you want to use a custom
+checkpoint path.
 
 ---
 
 ### Demo (with TTA)
 
-`demo_tta.py` pipeline: specify learning rate and number of iterations:
-
-Example:
+Run PRIMA inference with test-time adaptation:
 
 ```bash
-python demo_tta.py \
-  --checkpoint data/PRIMAS1/checkpoints/s1ckpt.ckpt \
-  --img_folder demo_data/ \
-  --out_folder demo_out_tta/ \
-  --tta_lr 1e-6 \
-  --tta_num_iters 30
+bash demo_tta.sh
 ```
 
-Outputs are written to `demo_out_tta/` (before/after TTA renders, keypoints, and optional meshes).
+Outputs are written to `demo_out_tta/` (before/after TTA renders, keypoints, and
+optional meshes). Edit `demo_tta.sh` if you want to change the checkpoint, TTA
+learning rate, or number of iterations.
 
 ---
 
@@ -123,12 +166,25 @@ browser:
 
 ```bash
 python app.py \
-  --checkpoint data/PRIMAS1/checkpoints/s1ckpt.ckpt \
+  --checkpoint data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt \
   --out_folder demo_out_tta_gradio/
 ```
 
 This starts a local Gradio app (by default on http://127.0.0.1:7860), where
 you can upload images and visualize PRIMA predictions and adaptation results.
+The `s1ckpt_inference.ckpt` checkpoint is downloaded automatically if missing.
+
+`app.py` picks a **demo profile** automatically:
+
+| | **Local** (`python app.py`) | **Hugging Face Space** |
+|--|--|--|
+| PRIMA device | GPU if available, else CPU | CPU only |
+| Detector | Detectron2 X-101-FPN | DeepLabCut SuperAnimal detector |
+| Default TTA iterations | 30 | 30 |
+| Save `.obj` meshes | on | off |
+| Preload checkpoint at startup | off | on |
+
+Override for testing: `PRIMA_DEMO_MODE=local` or `PRIMA_DEMO_MODE=space`.
 
 ---
 
@@ -155,7 +211,7 @@ Training outputs are written to `logs/train/runs/<exp_name>/`.
 ```bash
 python eval.py \
   --config data/PRIMAS1/.hydra/config.yaml \
-  --checkpoint data/PRIMAS1/checkpoints/s1ckpt.ckpt
+  --checkpoint data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt
 ```
 
 Common values for `--dataset` are controlled by:
@@ -171,6 +227,7 @@ This release builds on several open-source projects, including:
 - [BioCLIP](https://github.com/Imageomics/BioCLIP)
 - [AniMer](https://github.com/luoxue-star/AniMer)
 - [DeepLabCut](https://github.com/DeepLabCut/DeepLabCut)
+- [SAM3DB](https://github.com/facebookresearch/sam-3d-body)
 
 ---
 
@@ -182,7 +239,6 @@ If you use this code in your research, please cite our PRIMA paper.
 @misc{yu_prima,
   title={PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation},
   author={Xiaohang Yu and Ti Wang and Mackenzie Weygandt Mathis},
-  note={EPFL project page. Update publication year, venue, and links when available.}
 }
 ```
 
diff --git a/app.py b/app.py
index 15602d9..9654270 100644
--- a/app.py
+++ b/app.py
@@ -14,56 +14,265 @@
 
 1. Given an input image, run Detectron2 to detect animals.
 2. For each detected animal, run PRIMA for 3D pose/shape estimation.
-3. Run DeepLabCut SuperAnimal to obtain 2D keypoints.
-4. Map SuperAnimal 39 keypoints to the 26 PRIMA keypoints.
-5. Run test-time adaptation (TTA) with user-specified lr and iters.
-6. Render and save before/after TTA results and keypoint visualizations.
+3. Run the fine-tuned DeepLabCut SuperAnimal model to obtain PRIMA 26-keypoint
+   2D predictions.
+4. Run test-time adaptation (TTA) with user-specified lr and iters.
+5. Render and save before/after TTA results and keypoint visualizations.
 
 """
 
 import argparse
+import concurrent.futures
 import os
+import queue
+import sys
 import tempfile
+import time
+import traceback
+from dataclasses import dataclass
+from functools import lru_cache
 from types import SimpleNamespace
-from typing import List, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 from pathlib import Path
 
+# macOS: keep compute single-threaded and run inference on main thread.
+if sys.platform == "darwin":
+    os.environ.setdefault("OMP_NUM_THREADS", "1")
+
 import cv2
 import gradio as gr
 import numpy as np
 import torch
 import torch.utils.data
 
-import detectron2
-import detectron2.config
-import detectron2.engine
-from detectron2 import model_zoo
-
-from prima.models import load_prima
-from prima.utils import recursive_to
-from prima.datasets.vitdet_dataset import ViTDetDataset
-from prima.utils.renderer import Renderer
-
-# Reuse core utilities from the CLI demo_tta script
-from demo_tta import (
-    ANIMAL_COCO_IDS,
-    denorm_patch_to_rgb,
-    map_superanimal_to_prima,
-    run_superanimal_on_patch,
-    save_keypoint_vis,
-    tta_optimize,
+if sys.platform == "darwin":
+    torch.set_num_threads(1)
+
+# Repo-local minimal ``chumpy`` shim (see ``chumpy/__init__.py``) so SMAL pickles load
+# without installing the full chumpy package in Space builds.
+_REPO_ROOT = Path(__file__).resolve().parent
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from prima.utils.weights import (
+    DEFAULT_HF_REPO_ID,
+    resolve_prima_checkpoint_path,
 )
+from prima.utils.detection import select_animal_boxes
 
 
 # Default checkpoint path following README instructions
-DEFAULT_CHECKPOINT = "data/PRIMAS1/checkpoints/s1ckpt.ckpt"
+DEFAULT_CHECKPOINT = str(_REPO_ROOT / "data" / "PRIMAS1" / "checkpoints" / "s1ckpt_inference.ckpt")
+DEFAULT_HF_ASSET_REPO = DEFAULT_HF_REPO_ID
 
 # Output folder for rendered images/meshes and keypoints
 DEFAULT_OUT_FOLDER = "demo_out_tta_gradio"
+DEFAULT_SERVER_NAME = os.environ.get("PRIMA_GRADIO_HOST", "0.0.0.0")
+DEFAULT_SERVER_PORT = int(os.environ.get("PRIMA_GRADIO_PORT", "7860"))
+
+_D2_R50_CFG = "COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"
+_D2_R50_URL = (
+    "https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/"
+    "faster_rcnn_R_50_FPN_3x/137849458/model_final_280758.pkl"
+)
+_D2_X101_CFG = "COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"
+_D2_X101_URL = (
+    "https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/"
+    "faster_rcnn_X_101_32x8d_FPN_3x/139173657/model_final_68b088.pkl"
+)
+
+# Gradio example row: (image_rel, tta_lr, tta_iters, det_thresh, kp_thresh, side_view, render_depth, save_mesh)
+ExampleRow = Tuple[str, float, int, float, float, bool, bool, bool]
+
+
+@dataclass(frozen=True)
+class DemoProfile:
+    """Runtime settings for either the full local app or the lightweight HF Space demo."""
+
+    mode: str
+    prima_device: str  # "auto" (CUDA if available) or "cpu"
+    detectron_config_yaml: str
+    detectron_weights_url: str
+    detectron_device: str  # "auto" or "cpu"
+    default_tta_iters: int
+    max_tta_iters: int
+    default_save_mesh: bool
+    default_side_view: bool
+    default_render_depth: bool
+    preload_assets: bool
+    example_rows: Tuple[ExampleRow, ...]
+    description: str
+    interface_title: str
+
+    def resolve_prima_device(self) -> torch.device:
+        if self.prima_device == "cpu":
+            return torch.device("cpu")
+        return torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+    def resolve_detectron_device(self) -> str:
+        if self.detectron_device == "cpu":
+            return "cpu"
+        return "cuda" if torch.cuda.is_available() else "cpu"
+
+
+LOCAL_DEMO_PROFILE = DemoProfile(
+    mode="local",
+    prima_device="auto",
+    detectron_config_yaml=_D2_X101_CFG,
+    detectron_weights_url=_D2_X101_URL,
+    detectron_device="auto",
+    default_tta_iters=30,
+    max_tta_iters=100,
+    default_save_mesh=True,
+    default_side_view=False,
+    default_render_depth=False,
+    preload_assets=False,
+    example_rows=(
+        ("demo_data/000000015956_horse.png", 1e-6, 30, 0.7, 0.1, False, False, True),
+        ("demo_data/n02412080_12159.png", 1e-6, 30, 0.7, 0.1, False, False, True),
+        ("demo_data/000000315905_zebra.jpg", 1e-6, 30, 0.7, 0.1, False, False, True),
+        ("demo_data/beagle.jpg", 1e-6, 30, 0.7, 0.1, False, False, True),
+        ("demo_data/shepherd_hati.jpg", 1e-6, 30, 0.7, 0.1, False, False, True),
+    ),
+    description=(
+        "**Local demo** — full pipeline on your machine (GPU when available).\n\n"
+        "Detectron2 **X-101-FPN**, PRIMA mesh recovery, optional **DeepLabCut SuperAnimal + TTA**. "
+        "Set TTA iterations to **0** to skip adaptation. Outputs are saved under "
+        f"`{DEFAULT_OUT_FOLDER}`."
+    ),
+    interface_title=(
+        "PRIMA local demo (GPU/CPU) — detection, mesh recovery, optional TTA"
+    ),
+)
+
+SPACE_DEMO_PROFILE = DemoProfile(
+    mode="space",
+    prima_device="cpu",
+    detectron_config_yaml=_D2_R50_CFG,
+    detectron_weights_url=_D2_R50_URL,
+    detectron_device="cpu",
+    default_tta_iters=30,
+    max_tta_iters=30,
+    default_save_mesh=False,
+    default_side_view=False,
+    default_render_depth=False,
+    preload_assets=True,
+    example_rows=(
+        ("demo_data/000000015956_horse.png", 1e-6, 30, 0.7, 0.1, False, False, False),
+        ("demo_data/n02412080_12159.png", 1e-6, 30, 0.7, 0.1, False, False, False),
+        ("demo_data/000000315905_zebra.jpg", 1e-6, 30, 0.7, 0.1, False, False, False),
+        ("demo_data/beagle.jpg", 1e-6, 30, 0.7, 0.1, False, False, False),
+        ("demo_data/shepherd_hati.jpg", 1e-6, 30, 0.7, 0.1, False, False, False),
+    ),
+    description=(
+        "**Hugging Face Space (cpu-basic)** — lightweight demo: **CPU-only** PRIMA inference. "
+        "The Space build skips Detectron2 and uses the DeepLabCut SuperAnimal detector for animal "
+        "crops. TTA is optional (30 iterations by default, matching the local demo; set to 0 to "
+        "skip). Mesh `.obj` export is off by default to save time and disk."
+    ),
+    interface_title="PRIMA on Hugging Face — lightweight CPU demo",
+)
+
+
+def _is_truthy_env(var_name: str) -> bool:
+    return os.environ.get(var_name, "").strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _running_on_space() -> bool:
+    return bool(os.environ.get("SPACE_ID") or os.environ.get("HF_SPACE_ID"))
+
+
+@lru_cache(maxsize=1)
+def get_demo_profile() -> DemoProfile:
+    """Select local vs Space profile. Override with ``PRIMA_DEMO_MODE=local|space``."""
+    override = os.environ.get("PRIMA_DEMO_MODE", "").strip().lower()
+    if override == "local":
+        return LOCAL_DEMO_PROFILE
+    if override == "space":
+        return SPACE_DEMO_PROFILE
+    return SPACE_DEMO_PROFILE if _running_on_space() else LOCAL_DEMO_PROFILE
+
+
+def _gradio_examples_for_interface(profile: DemoProfile) -> List[List]:
+    """Gradio prefetches example media at startup (paths must exist beside ``app.py``)."""
+    if _is_truthy_env("PRIMA_DISABLE_GRADIO_EXAMPLES"):
+        return []
+    rows: List[List] = []
+    for rel, *rest in profile.example_rows:
+        p = _REPO_ROOT / rel
+        if p.is_file():
+            rows.append([str(p), *rest])
+    return rows
+
+
+def _should_use_gradio_queue(profile: DemoProfile) -> bool:
+    """Whether to enable Gradio's background queue.
+
+    On macOS local, the queue runs handlers in worker threads; PyRender and
+    pyglet/AppKit then crash with "setting the main menu on a non-main thread".
+    """
+    if _is_truthy_env("PRIMA_GRADIO_NO_QUEUE"):
+        return False
+    if _is_truthy_env("PRIMA_GRADIO_QUEUE"):
+        return True
+    return not (sys.platform == "darwin" and profile.mode == "local")
+
+
+def _warmup_runtime_cache(checkpoint_path: str, profile: DemoProfile) -> Dict[str, Any]:
+    """Load model + DLC on the main thread (recommended for macOS local)."""
+    print("[startup] Preloading DeepLabCut on main thread…")
+    _deeplabcut_available()
+    print("[startup] Loading PRIMA + Detectron2 on main thread (first run can take several minutes)…")
+    model, model_cfg, renderer, cam_crop_to_full_fn, device, detector = _load_model_and_detector_for_demo(
+        checkpoint_path, profile
+    )
+    return {
+        "model": model,
+        "model_cfg": model_cfg,
+        "renderer": renderer,
+        "cam_crop_to_full_fn": cam_crop_to_full_fn,
+        "device": device,
+        "detector": detector,
+    }
+
+
+def _should_preload_assets(profile: DemoProfile) -> bool:
+    preload_env = os.environ.get("PRIMA_PRELOAD_ASSETS")
+    if preload_env is not None:
+        return _is_truthy_env("PRIMA_PRELOAD_ASSETS")
+    return profile.preload_assets
+
+def _deeplabcut_available() -> bool:
+    try:
+        from deeplabcut.pose_estimation_pytorch.apis import superanimal_analyze_images  # noqa: F401
+
+        return True
+    except Exception:
+        return False
+
+
+def _preload_assets_once(checkpoint_path: str) -> None:
+    print("[startup] Ensuring demo assets from Hugging Face Hub...")
+    resolve_prima_checkpoint_path(
+        checkpoint_path,
+        data_dir=_REPO_ROOT / "data",
+        auto_download=True,
+        hf_repo_id=os.environ.get("PRIMA_HF_REPO_ID", DEFAULT_HF_ASSET_REPO),
+    )
+    print("[startup] Asset preload complete.")
 
 
 def _load_prima_model(checkpoint_path: str = DEFAULT_CHECKPOINT):
     """Load PRIMA model and renderer once for the Gradio app."""
+    from prima.models import load_prima
+    from prima.utils.renderer import Renderer, cam_crop_to_full
+
+    checkpoint_path = resolve_prima_checkpoint_path(
+        checkpoint_path,
+        data_dir=_REPO_ROOT / "data",
+        auto_download=True,
+        hf_repo_id=os.environ.get("PRIMA_HF_REPO_ID", DEFAULT_HF_ASSET_REPO),
+    )
     checkpoint = Path(checkpoint_path)
     cfg_path = checkpoint.parent.parent / ".hydra" / "config.yaml"
     if not checkpoint.exists():
@@ -75,36 +284,155 @@ def _load_prima_model(checkpoint_path: str = DEFAULT_CHECKPOINT):
             f"Missing model config: {cfg_path}. Ensure the full checkpoint folder layout from README is present."
         )
 
+    profile = get_demo_profile()
     model, model_cfg = load_prima(checkpoint_path)
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    device = profile.resolve_prima_device()
     model = model.to(device)
     model.eval()
 
     renderer = Renderer(model_cfg, faces=model.smal.faces)
-    return model, model_cfg, renderer, device
-
-
-def _build_detector():
-    """Build Detectron2 animal detector (same config as demo_tta/demo.py)."""
+    return model, model_cfg, renderer, cam_crop_to_full, device
+
+
+def _build_detector(profile: Optional[DemoProfile] = None):
+    """Build Detectron2 animal detector (profile selects X-101+GPU locally vs R50+CPU on Space)."""
+    try:
+        import detectron2.config
+        import detectron2.engine
+        from detectron2 import model_zoo
+    except Exception as e:
+        print(f"[warn] Detectron2 unavailable ({type(e).__name__}: {e}); using SuperAnimal detector fallback.")
+        return None
+
+    if profile is None:
+        profile = get_demo_profile()
+    config_yaml = profile.detectron_config_yaml
+    weights = profile.detectron_weights_url
+    device_str = profile.resolve_detectron_device()
+    print(f"[detectron2] mode={profile.mode} config={config_yaml} device={device_str}")
 
     cfg = detectron2.config.get_cfg()
-    cfg.merge_from_file(
-        model_zoo.get_config_file("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml")
-    )
+    cfg.merge_from_file(model_zoo.get_config_file(config_yaml))
     cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
-    cfg.MODEL.WEIGHTS = (
-        "https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/"
-        "faster_rcnn_X_101_32x8d_FPN_3x/139173657/model_final_68b088.pkl"
-    )
+    cfg.MODEL.WEIGHTS = weights
+    cfg.MODEL.DEVICE = device_str
     detector = detectron2.engine.DefaultPredictor(cfg)
     return detector
 
+
+def _filter_superanimal_boxes(
+    payload: Dict[str, Any],
+    det_thresh: float,
+    img_shape: Tuple[int, int],
+) -> Optional[np.ndarray]:
+    boxes = payload.get("bboxes")
+    if boxes is None:
+        return None
+
+    boxes = np.asarray(boxes, dtype=np.float32).reshape(-1, 4)
+    if len(boxes) == 0:
+        return None
+
+    scores = payload.get("bbox_scores")
+    if scores is None:
+        scores = np.ones((len(boxes),), dtype=np.float32)
+    scores = np.asarray(scores, dtype=np.float32).reshape(-1)
+    if len(scores) != len(boxes):
+        scores = np.ones((len(boxes),), dtype=np.float32)
+
+    h, w = img_shape
+    valid = (
+        (scores > float(det_thresh))
+        & np.isfinite(boxes).all(axis=1)
+        & (boxes[:, 2] > 0.0)
+        & (boxes[:, 3] > 0.0)
+    )
+    if not np.any(valid):
+        return None
+
+    xywh = boxes[valid]
+    scores = scores[valid]
+    boxes = xywh.copy()
+    boxes[:, 2] = xywh[:, 0] + xywh[:, 2]
+    boxes[:, 3] = xywh[:, 1] + xywh[:, 3]
+    boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0.0, float(max(1, w - 1)))
+    boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0.0, float(max(1, h - 1)))
+    valid_size = (boxes[:, 2] > boxes[:, 0]) & (boxes[:, 3] > boxes[:, 1])
+    boxes = boxes[valid_size]
+    scores = scores[valid_size]
+    if len(boxes) == 0:
+        return None
+    order = np.argsort(scores)[::-1]
+    return boxes[order].astype(np.float32, copy=False)
+
+
+def _detect_superanimal_boxes(img_rgb: np.ndarray, det_thresh: float) -> Optional[np.ndarray]:
+    try:
+        from deeplabcut.pose_estimation_pytorch.apis import superanimal_analyze_images
+    except Exception as e:
+        print(f"[warn] DeepLabCut SuperAnimal unavailable ({type(e).__name__}: {e}); no fallback bbox.")
+        return None
+
+    with tempfile.TemporaryDirectory(prefix="sa_detect_") as tmp_dir:
+        img_path = os.path.join(tmp_dir, "image.png")
+        cv2.imwrite(img_path, cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR))
+
+        dlc_device = "cuda" if torch.cuda.is_available() else "cpu"
+        preds = superanimal_analyze_images(
+            superanimal_name=SUPER_ANIMAL_ARGS.superanimal_name,
+            model_name=SUPER_ANIMAL_ARGS.superanimal_model_name,
+            detector_name=SUPER_ANIMAL_ARGS.superanimal_detector_name,
+            images=img_path,
+            max_individuals=SUPER_ANIMAL_ARGS.superanimal_max_individuals,
+            out_folder=tmp_dir,
+            progress_bar=False,
+            device=dlc_device,
+            pose_threshold=0.1,
+            bbox_threshold=float(det_thresh),
+            plot_skeleton=False,
+        )
+
+        payload = preds.get(img_path)
+        if payload is None:
+            return None
+        return _filter_superanimal_boxes(payload, det_thresh, img_rgb.shape[:2])
+
+
+def _load_model_and_detector_for_demo(checkpoint_path: str, profile: DemoProfile):
+    """Load PRIMA and Detectron2 once for the Gradio session (main thread only)."""
+    model, model_cfg, renderer, cam_crop_to_full_fn, device = _load_prima_model(checkpoint_path)
+    detector = _build_detector(profile)
+    return model, model_cfg, renderer, cam_crop_to_full_fn, device, detector
+
+
+def _detect_animal_boxes(
+    detector,
+    img_bgr: np.ndarray,
+    det_thresh: float,
+) -> Optional[np.ndarray]:
+    """Return Nx4 XYXY boxes or None if no animal detections."""
+    if detector is None:
+        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+        return _detect_superanimal_boxes(img_rgb, det_thresh)
+
+    det_out = detector(img_bgr)
+    det_instances = det_out["instances"]
+    boxes, suppressed = select_animal_boxes(det_instances, score_threshold=float(det_thresh))
+    if suppressed > 0:
+        print(f"[INFO] Suppressed {suppressed} duplicate animal detection(s)")
+    if len(boxes) == 0:
+        return None
+    return boxes
+
+
 # SuperAnimal defaults (same as in demo_tta parser)
 SUPER_ANIMAL_ARGS = SimpleNamespace(
     superanimal_name="superanimal_quadruped",
     superanimal_model_name="hrnet_w32",
     superanimal_detector_name="fasterrcnn_resnet50_fpn_v2",
     superanimal_max_individuals=1,
+    saved_2d_model_path="",
+    pytorch_config_2d_path=str(_REPO_ROOT / "configs" / "sa_finetune_hrnet_w32.yaml"),
 )
 
 
@@ -112,6 +440,7 @@ def _collect_animal_results(
     model,
     model_cfg,
     renderer,
+    cam_crop_to_full_fn,
     device,
     detector,
     out_folder: str,
@@ -121,7 +450,10 @@ def _collect_animal_results(
     det_thresh: float,
     kp_conf_thresh: float,
     side_view: bool,
+    render_depth: bool,
     save_mesh: bool,
+    boxes: Optional[np.ndarray] = None,
+    progress_callback: Optional[Callable[[str], None]] = None,
 ) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray], str | None, str | None]:
     """Run detection + PRIMA + SuperAnimal + TTA on a single RGB image.
 
@@ -132,22 +464,35 @@ def _collect_animal_results(
         first_before_mesh: path to first animal's before-TTA mesh (.obj) or None
         first_after_mesh: path to first animal's after-TTA mesh (.obj) or None
     """
+    from prima.utils import recursive_to
+    from prima.datasets.vitdet_dataset import ViTDetDataset
+    from demo_tta import (
+        denorm_patch_to_rgb,
+        resolve_sa_weights_path,
+        run_superanimal_on_patch,
+        save_keypoint_vis,
+        tta_optimize,
+    )
 
-    # Detect animals
-    img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
-    det_out = detector(img_bgr)
-    det_instances = det_out["instances"]
+    def report(message: str) -> None:
+        if progress_callback is not None:
+            progress_callback(message)
 
-    valid_idx = [
-        i
-        for i, (c, s) in enumerate(zip(det_instances.pred_classes, det_instances.scores))
-        if (int(c) in ANIMAL_COCO_IDS) and (float(s) > float(det_thresh))
-    ]
-    if len(valid_idx) == 0:
-        return [], [], [], None, None
+    if int(tta_num_iters) > 0 and not SUPER_ANIMAL_ARGS.saved_2d_model_path:
+        report("Resolving SuperAnimal weights...")
+        SUPER_ANIMAL_ARGS.saved_2d_model_path = resolve_sa_weights_path("")
 
-    boxes = det_instances.pred_boxes.tensor[valid_idx].cpu().numpy()
+    img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
+    if boxes is None:
+        if detector is None:
+            report("Detectron2 unavailable; detecting animals with SuperAnimal...")
+        else:
+            report("Detecting animals with Detectron2...")
+        boxes = _detect_animal_boxes(detector, img_bgr, det_thresh)
+    if boxes is None:
+        return [], [], [], None, None
 
+    report(f"Detected {len(boxes)} animal(s). Preparing crops...")
     dataset = ViTDetDataset(model_cfg, img_bgr, boxes)
     dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)
 
@@ -159,9 +504,11 @@ def _collect_animal_results(
 
     img_token = next(tempfile._get_candidate_names())
 
-    for batch in dataloader:
+    total_batches = len(dataloader)
+    for batch_idx, batch in enumerate(dataloader, start=1):
         batch = recursive_to(batch, device)
 
+        report(f"Animal {batch_idx}/{total_batches}: running PRIMA...")
         with torch.no_grad():
             out_before = model(batch)
 
@@ -171,8 +518,10 @@ def _collect_animal_results(
         img_fn = f"{img_token}"
         from demo_tta import render_and_save  # imported lazily to avoid circular issues
 
+        report(f"Animal {batch_idx}/{total_batches}: rendering before TTA...")
         render_and_save(
             renderer,
+            cam_crop_to_full_fn,
             out_before,
             batch,
             img_fn,
@@ -181,6 +530,7 @@ def _collect_animal_results(
             suffix="before_tta",
             side_view=side_view,
             save_mesh=save_mesh,
+            render_depth=render_depth,
         )
 
         before_png_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_before_tta.png")
@@ -195,8 +545,10 @@ def _collect_animal_results(
                 before_mesh_paths.append(before_obj_path)
 
         if int(tta_num_iters) <= 0:
+            report(f"Animal {batch_idx}/{total_batches}: rendering final output...")
             render_and_save(
                 renderer,
+                cam_crop_to_full_fn,
                 out_before,
                 batch,
                 img_fn,
@@ -205,6 +557,7 @@ def _collect_animal_results(
                 suffix="after_tta",
                 side_view=side_view,
                 save_mesh=save_mesh,
+                render_depth=render_depth,
             )
 
             after_png_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_after_tta.png")
@@ -220,6 +573,7 @@ def _collect_animal_results(
             continue
 
         # Prepare patch for SuperAnimal
+        report(f"Animal {batch_idx}/{total_batches}: running SuperAnimal keypoints...")
         patch_rgb = denorm_patch_to_rgb(batch["img"][0])
         with tempfile.TemporaryDirectory(prefix=f"dlc_{img_fn}_{animal_id}_") as tmp_dir:
             bodyparts_xyc = run_superanimal_on_patch(patch_rgb, SUPER_ANIMAL_ARGS, tmp_dir)
@@ -228,14 +582,14 @@ def _collect_animal_results(
             # No keypoints => skip TTA for this animal
             continue
 
-        mapped_xyc = map_superanimal_to_prima(bodyparts_xyc)
-        mapped_xyc[mapped_xyc[:, 2] < float(kp_conf_thresh), 2] = 0.0
+        kpts_xyc = bodyparts_xyc
+        kpts_xyc[kpts_xyc[:, 2] < float(kp_conf_thresh), 2] = 0.0
 
         # Save keypoint visualization and npy
         kpt_png_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_prima26_kpts.png")
-        save_keypoint_vis(patch_rgb, mapped_xyc, kpt_png_path)
+        save_keypoint_vis(patch_rgb, kpts_xyc, kpt_png_path)
         npy_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_prima26_kpts.npy")
-        np.save(npy_path, mapped_xyc)
+        np.save(npy_path, kpts_xyc)
 
         if os.path.exists(kpt_png_path):
             kpt_bgr = cv2.imread(kpt_png_path)
@@ -244,12 +598,13 @@ def _collect_animal_results(
 
         # Normalize keypoints to [-0.5, 0.5] as in demo_tta
         patch_h, patch_w = patch_rgb.shape[:2]
-        mapped_norm = mapped_xyc.copy()
-        mapped_norm[:, 0] = mapped_norm[:, 0] / float(patch_w) - 0.5
-        mapped_norm[:, 1] = mapped_norm[:, 1] / float(patch_h) - 0.5
-        gt_kpts_norm = torch.from_numpy(mapped_norm[None]).to(device=device, dtype=batch["img"].dtype)
+        kpts_norm = kpts_xyc.copy()
+        kpts_norm[:, 0] = kpts_norm[:, 0] / float(patch_w) - 0.5
+        kpts_norm[:, 1] = kpts_norm[:, 1] / float(patch_h) - 0.5
+        gt_kpts_norm = torch.from_numpy(kpts_norm[None]).to(device=device, dtype=batch["img"].dtype)
 
         # Run TTA
+        report(f"Animal {batch_idx}/{total_batches}: running TTA ({int(tta_num_iters)} iterations)...")
         out_after = tta_optimize(
             model,
             batch,
@@ -258,8 +613,10 @@ def _collect_animal_results(
             lr=float(tta_lr),
         )
 
+        report(f"Animal {batch_idx}/{total_batches}: rendering after TTA...")
         render_and_save(
             renderer,
+            cam_crop_to_full_fn,
             out_after,
             batch,
             img_fn,
@@ -268,6 +625,7 @@ def _collect_animal_results(
             suffix="after_tta",
             side_view=side_view,
             save_mesh=save_mesh,
+            render_depth=render_depth,
         )
 
         after_png_path = os.path.join(out_folder, f"{img_fn}_{animal_id}_after_tta.png")
@@ -284,13 +642,33 @@ def _collect_animal_results(
     first_before_mesh = before_mesh_paths[0] if before_mesh_paths else None
     first_after_mesh = after_mesh_paths[0] if after_mesh_paths else None
 
+    report("Collecting outputs...")
     return before_imgs, after_imgs, kpt_imgs, first_before_mesh, first_after_mesh
 
 
-def build_demo(checkpoint_path: str = DEFAULT_CHECKPOINT, out_folder: str = DEFAULT_OUT_FOLDER) -> gr.Interface:
+def build_demo(
+    checkpoint_path: str = DEFAULT_CHECKPOINT,
+    out_folder: str = DEFAULT_OUT_FOLDER,
+    runtime_cache: Optional[Dict[str, Any]] = None,
+) -> gr.Interface:
+    profile = get_demo_profile()
+    print(
+        f"[demo] profile={profile.mode} prima={profile.resolve_prima_device()} "
+        f"detectron={profile.detectron_config_yaml} d2_device={profile.resolve_detectron_device()}"
+    )
+    if _should_use_gradio_queue(profile):
+        print("[demo] Gradio queue enabled (background worker threads).")
+    else:
+        print("[demo] Gradio queue disabled (inference runs on main thread; required on macOS local).")
     os.makedirs(out_folder, exist_ok=True)
-    model, model_cfg, renderer, device = _load_prima_model(checkpoint_path)
-    detector = _build_detector()
+    runtime_cache = runtime_cache or {
+        "model": None,
+        "model_cfg": None,
+        "renderer": None,
+        "cam_crop_to_full_fn": None,
+        "device": None,
+        "detector": None,
+    }
 
     def gradio_inference(
         image: np.ndarray,
@@ -299,39 +677,157 @@ def gradio_inference(
         det_thresh: float,
         kp_conf_thresh: float,
         side_view: bool,
+        render_depth: bool,
         save_mesh: bool,
     ):
-        """Wrapper for Gradio. ``image`` is an RGB numpy array."""
+        """Wrapper for Gradio. ``image`` is an RGB numpy array.
+
+        Yields intermediate status so long first-run (Hub downloads + model load)
+        and long inference do not hit silent client/proxy WebSocket timeouts.
+        """
 
         if image is None:
-            return [], [], [], None, None
+            yield None, None, None, "No image provided."
+            return
+
+        if int(tta_num_iters) > 0 and not _deeplabcut_available():
+            yield (
+                None,
+                None,
+                None,
+                "DeepLabCut is not installed. Set **TTA iterations** to **0** for PRIMA-only inference, "
+                "or install `deeplabcut` (see README / requirements.txt).",
+            )
+            return
 
         if image.dtype != np.uint8:
             img_rgb = np.clip(image, 0, 255).astype(np.uint8)
         else:
             img_rgb = image
 
-        before_imgs, after_imgs, kpt_imgs, mesh_before, mesh_after = _collect_animal_results(
-            model,
-            model_cfg,
-            renderer,
-            device,
-            detector,
-            out_folder,
-            img_rgb,
-            tta_lr=tta_lr,
-            tta_num_iters=tta_num_iters,
-            det_thresh=det_thresh,
-            kp_conf_thresh=kp_conf_thresh,
-            side_view=side_view,
-            save_mesh=save_mesh,
-        )
+        yield None, None, None, "Queued; preparing run…"
 
-        return before_imgs, after_imgs, kpt_imgs, mesh_before, mesh_after
+        if runtime_cache["model"] is None:
+            yield (
+                None,
+                None,
+                None,
+                "First run: downloading demo assets from Hugging Face (large checkpoint) "
+                "and loading the model. This can take many minutes.",
+            )
+            try:
+                model, model_cfg, renderer, cam_crop_to_full_fn, device, detector = _load_model_and_detector_for_demo(
+                    checkpoint_path, profile
+                )
+            except Exception:
+                yield None, None, None, f"Model initialization failed:\n{traceback.format_exc()}"
+                return
+            runtime_cache["model"] = model
+            runtime_cache["model_cfg"] = model_cfg
+            runtime_cache["renderer"] = renderer
+            runtime_cache["cam_crop_to_full_fn"] = cam_crop_to_full_fn
+            runtime_cache["device"] = device
+            runtime_cache["detector"] = detector
+            yield None, None, None, "Model loaded."
+
+        try:
+            yield None, None, None, "Running animal detection…"
+            img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
+            boxes = _detect_animal_boxes(runtime_cache["detector"], img_bgr, det_thresh)
+            if boxes is None:
+                yield (
+                    None,
+                    None,
+                    None,
+                    "No animal detected. Try lowering the detection threshold or another image.",
+                )
+                return
+            yield (
+                None,
+                None,
+                None,
+                f"Detected {len(boxes)} animal region(s). Running PRIMA (+ SuperAnimal/TTA if enabled)...",
+            )
+
+            def run_collect(progress_callback: Optional[Callable[[str], None]] = None):
+                return _collect_animal_results(
+                    runtime_cache["model"],
+                    runtime_cache["model_cfg"],
+                    runtime_cache["renderer"],
+                    runtime_cache["cam_crop_to_full_fn"],
+                    runtime_cache["device"],
+                    runtime_cache["detector"],
+                    out_folder,
+                    img_rgb,
+                    tta_lr=tta_lr,
+                    tta_num_iters=tta_num_iters,
+                    det_thresh=det_thresh,
+                    kp_conf_thresh=kp_conf_thresh,
+                    side_view=side_view,
+                    render_depth=render_depth,
+                    save_mesh=save_mesh,
+                    boxes=boxes,
+                    progress_callback=progress_callback,
+                )
+
+            if _should_use_gradio_queue(profile):
+                stage_updates: queue.Queue[str] = queue.Queue()
+
+                def report_stage(message: str) -> None:
+                    stage_updates.put(message)
+
+                with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
+                    fut = pool.submit(
+                        run_collect,
+                        report_stage,
+                    )
+                    t0 = time.monotonic()
+                    latest_stage = "Starting inference..."
+                    while True:
+                        while True:
+                            try:
+                                latest_stage = stage_updates.get_nowait()
+                            except queue.Empty:
+                                break
+                            else:
+                                elapsed = int(time.monotonic() - t0)
+                                yield None, None, None, f"{latest_stage}\nElapsed: {elapsed}s"
+                        try:
+                            before_imgs, after_imgs, kpt_imgs, mesh_before, mesh_after = fut.result(
+                                timeout=1.0
+                            )
+                            break
+                        except concurrent.futures.TimeoutError:
+                            elapsed = int(time.monotonic() - t0)
+                            yield None, None, None, (
+                                f"{latest_stage}\n"
+                                f"Elapsed: {elapsed}s\n"
+                                "CPU inference can take several minutes."
+                            )
+            else:
+                before_imgs, after_imgs, kpt_imgs, mesh_before, mesh_after = run_collect()
+        except Exception:
+            yield None, None, None, f"Inference failed:\n{traceback.format_exc()}"
+            return
+
+        first_before = before_imgs[0] if before_imgs else None
+        first_after = after_imgs[0] if after_imgs else None
+        first_kpts = kpt_imgs[0] if kpt_imgs else None
+        if first_before is None and first_after is None:
+            yield (
+                None,
+                None,
+                None,
+                "No output generated. Try an image with a clearly visible quadruped.",
+            )
+            return
+        yield first_before, first_after, first_kpts, "OK"
 
-    return gr.Interface(
+    _gradio_examples = _gradio_examples_for_interface(profile)
+    _iface_kw = dict(
         fn=gradio_inference,
         analytics_enabled=False,
+        cache_examples=False,
         inputs=[
             gr.Image(
                 label="Input image",
@@ -348,8 +844,8 @@ def gradio_inference(
             gr.Slider(
                 label="TTA iterations",
                 minimum=0,
-                maximum=100,
-                value=30,
+                maximum=profile.max_tta_iters,
+                value=profile.default_tta_iters,
                 step=1,
                 info="Set to 0 to disable TTA and reuse the initial PRIMA prediction.",
             ),
@@ -367,73 +863,25 @@ def gradio_inference(
                 value=0.1,
                 step=0.05,
             ),
-            gr.Checkbox(label="Render side view", value=False),
-            gr.Checkbox(label="Save meshes (.obj)", value=True),
+            gr.Checkbox(label="Render side view", value=profile.default_side_view),
+            gr.Checkbox(label="Render depth map", value=profile.default_render_depth),
+            gr.Checkbox(label="Save meshes (.obj)", value=profile.default_save_mesh),
         ],
         outputs=[
-            gr.Gallery(label="Before TTA (all animals)"),
-            gr.Gallery(label="After TTA (all animals)"),
-            gr.Gallery(label="PRIMA 26 keypoints"),
-            gr.Model3D(label="First animal mesh before TTA"),
-            gr.Model3D(label="First animal mesh after TTA"),
-        ],
-        title="PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation",
-        description=(
-            "Upload an animal image. The demo runs Detectron2 for animal detection, "
-            "PRIMA for 3D pose/shape, DeepLabCut SuperAnimal for 2D keypoints, and "
-            "test-time adaptation (TTA) with configurable learning rate and iterations. "
-            "Set TTA iterations to 0 to disable adaptation.\n\n"
-            "Results (PNG/OBJ and 26-keypoint visualizations) are saved under "
-            f"'{out_folder}'."
-        ),
-        examples=[
-            [
-                "demo_data/000000015956_horse.png",
-                1e-6,
-                30,
-                0.7,
-                0.1,
-                False,
-                True,
-            ],
-            [
-                "demo_data/n02412080_12159.png",
-                1e-6,
-                30,
-                0.7,
-                0.1,
-                False,
-                True,
-            ],
-            [
-                "demo_data/000000315905_zebra.jpg",
-                1e-6,
-                30,
-                0.7,
-                0.1,
-                False,
-                True,
-            ],
-            [
-                "demo_data/beagle.jpg",
-                1e-6,
-                0,
-                0.7,
-                0.1,
-                False,
-                True,
-            ],
-            [
-                "demo_data/shepherd_hati.jpg",
-                1e-6,
-                0,
-                0.7,
-                0.1,
-                False,
-                True,
-            ],
+            gr.Image(label="Before TTA"),
+            gr.Image(label="After TTA"),
+            gr.Image(label="PRIMA 26 keypoints"),
+            gr.Textbox(label="Status / Traceback", lines=12),
         ],
+        title=profile.interface_title,
+        description=profile.description,
     )
+    if _gradio_examples:
+        _iface_kw["examples"] = _gradio_examples
+    demo = gr.Interface(**_iface_kw)
+    if _should_use_gradio_queue(profile):
+        demo.queue(max_size=8, default_concurrency_limit=1)
+    return demo
 
 
 def parse_args() -> argparse.Namespace:
@@ -450,10 +898,41 @@ def parse_args() -> argparse.Namespace:
         default=DEFAULT_OUT_FOLDER,
         help="Folder used to save rendered outputs and meshes",
     )
+    parser.add_argument(
+        "--server_name",
+        type=str,
+        default=DEFAULT_SERVER_NAME,
+        help="Host/interface used by Gradio. Use 0.0.0.0 for Run:AI port-forward.",
+    )
+    parser.add_argument(
+        "--server_port",
+        type=int,
+        default=DEFAULT_SERVER_PORT,
+        help="Port used by Gradio.",
+    )
     return parser.parse_args()
 
 
 if __name__ == "__main__":
     args = parse_args()
-    demo = build_demo(checkpoint_path=args.checkpoint, out_folder=args.out_folder)
-    demo.launch()
+    profile = get_demo_profile()
+    if _should_preload_assets(profile):
+        _preload_assets_once(args.checkpoint)
+    runtime_cache: Optional[Dict[str, Any]] = None
+    if (
+        sys.platform == "darwin"
+        and profile.mode == "local"
+        and _is_truthy_env("PRIMA_WARMUP")
+    ):
+        runtime_cache = _warmup_runtime_cache(args.checkpoint, profile)
+    demo = build_demo(
+        checkpoint_path=args.checkpoint,
+        out_folder=args.out_folder,
+        runtime_cache=runtime_cache,
+    )
+    demo.launch(
+        inbrowser=False,
+        ssr_mode=False,
+        server_name=args.server_name,
+        server_port=args.server_port,
+    )
diff --git a/chumpy/__init__.py b/chumpy/__init__.py
new file mode 100644
index 0000000..bb5b170
--- /dev/null
+++ b/chumpy/__init__.py
@@ -0,0 +1,16 @@
+from __future__ import annotations
+"""
+PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation
+
+Official implementation of the paper:
+"PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation"
+by Xiaohang Yu, Ti Wang, and Mackenzie Weygandt Mathis
+Licensed under a modified MIT license
+"""
+
+
+"""Minimal ``chumpy`` compatibility for unpickling legacy SMAL model configs."""
+
+from .ch import Ch, ChArray, materialize
+
+__all__ = ["Ch", "ChArray", "materialize"]
diff --git a/chumpy/ch.py b/chumpy/ch.py
new file mode 100644
index 0000000..7175103
--- /dev/null
+++ b/chumpy/ch.py
@@ -0,0 +1,90 @@
+from __future__ import annotations
+"""
+PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation
+
+Official implementation of the paper:
+"PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation"
+by Xiaohang Yu, Ti Wang, and Mackenzie Weygandt Mathis
+Licensed under a modified MIT license
+"""
+
+
+"""``chumpy.ch`` namespace expected by legacy SMAL pickles."""
+
+import numpy as np
+
+
+class Ch:
+    """Minimal stand-in for ``chumpy.ch.Ch`` (unpickling only)."""
+
+    def __init__(self, *args, **kwargs):
+        self._data = None
+        if args:
+            self._data = np.asarray(args[0])
+
+    def _resolve(self) -> np.ndarray:
+        # Real chumpy Ch instances store the underlying ndarray on attribute ``x``;
+        # legacy pickles unpickle by restoring ``__dict__`` without calling ``__init__``,
+        # so try common attribute names before falling back to ``_data``.
+        r = self.__dict__.get("r")
+        if isinstance(r, np.ndarray):
+            return np.asarray(r)
+        for attr in ("x", "_x", "_data"):
+            val = self.__dict__.get(attr)
+            if val is not None:
+                return np.asarray(val)
+        data = self.__dict__.get("_data")
+        if data is not None:
+            return np.asarray(data)
+        return np.zeros((), dtype=np.float32)
+
+    def r(self) -> np.ndarray:
+        """Match real chumpy API (``ch.r()`` returns the underlying array)."""
+        return self._resolve()
+
+    def __array__(self, dtype=None):
+        arr = self._resolve()
+        if dtype is not None:
+            arr = arr.astype(dtype, copy=False)
+        return arr
+
+
+class ChArray(np.ndarray):
+    """Minimal stand-in for ``chumpy.ch.ChArray``."""
+
+
+def _unwrap_ch(value, dtype=np.float32) -> np.ndarray:
+    """Resolve a chumpy ``Ch`` (method, property, or unpickled ``r``/``x`` attrs) to ndarray."""
+    if isinstance(value, Ch):
+        return np.asarray(value._resolve(), dtype=dtype)
+    r = getattr(value, "r", None)
+    if isinstance(r, np.ndarray):
+        return np.asarray(r, dtype=dtype)
+    if callable(r):
+        return np.asarray(r(), dtype=dtype)
+    for attr in ("x", "_x", "_data"):
+        val = getattr(value, attr, None)
+        if val is not None:
+            return np.asarray(val, dtype=dtype)
+    raise TypeError(f"Cannot materialize chumpy-like object: {type(value)!r}")
+
+
+def materialize(value, dtype=np.float32) -> np.ndarray:
+    """Recursively unwrap ``Ch`` / object arrays from legacy SMAL pickles."""
+    if isinstance(value, Ch) or (
+        type(value).__name__ == "Ch"
+        and hasattr(value, "r")
+        and not isinstance(value, (np.ndarray, list, tuple, dict, str, bytes))
+    ):
+        return _unwrap_ch(value, dtype=dtype)
+    if isinstance(value, np.ndarray):
+        if value.dtype == object:
+            flat = [materialize(x, dtype=dtype) for x in value.ravel()]
+            return np.stack(flat).reshape(value.shape)
+        return np.asarray(value, dtype=dtype)
+    if isinstance(value, (list, tuple)):
+        return np.asarray([materialize(x, dtype=dtype) for x in value], dtype=dtype)
+    return np.asarray(value, dtype=dtype)
+
+
+__all__ = ["Ch", "ChArray", "materialize"]
diff --git a/configs/sa_finetune_hrnet_w32.yaml b/configs/sa_finetune_hrnet_w32.yaml
new file mode 100644
index 0000000..83dd64e
--- /dev/null
+++ b/configs/sa_finetune_hrnet_w32.yaml
@@ -0,0 +1,220 @@
+# DeepLabCut pytorch_config for the PRIMA TTA 2D pose model:
+# SuperAnimal-Quadruped HRNet-w32 backbone fine-tuned on Animal3D, with
+# the heatmap head re-trained for the 26-joint Animal3D / PRIMA layout.
+#
+# Used by demo_tta.py via DLC's `superanimal_analyze_images(...,
+# customized_model_config=<this yaml>, customized_pose_checkpoint=<your
+# fine-tuned .pt>)`. Only the pose model is fine-tuned; the bounding-box
+# detector (Faster R-CNN) is the stock SuperAnimal-Quadruped one
+# resolved by DLC at runtime.
+data:
+  bbox_margin: 20
+  colormode: RGB
+  inference:
+    normalize_images: true
+    top_down_crop:
+      width: 256
+      height: 256
+    auto_padding:
+      pad_width_divisor: 32
+      pad_height_divisor: 32
+  train:
+    affine:
+      p: 0.5
+      rotation: 30
+      scaling:
+      - 1.0
+      - 1.0
+      translation: 0
+    gaussian_noise: 12.75
+    motion_blur: true
+    normalize_images: true
+    top_down_crop:
+      width: 256
+      height: 256
+    auto_padding:
+      pad_width_divisor: 32
+      pad_height_divisor: 32
+detector:
+  data:
+    colormode: RGB
+    inference:
+      normalize_images: true
+    train:
+      affine:
+        p: 0.5
+        rotation: 30
+        scaling:
+        - 1.0
+        - 1.0
+        translation: 40
+      collate:
+        type: ResizeFromDataSizeCollate
+        min_scale: 0.4
+        max_scale: 1.0
+        min_short_side: 128
+        max_short_side: 1152
+        multiple_of: 32
+        to_square: false
+      hflip: true
+      normalize_images: true
+  device: auto
+  model:
+    type: FasterRCNN
+    freeze_bn_stats: true
+    freeze_bn_weights: false
+    variant: fasterrcnn_resnet50_fpn_v2
+  runner:
+    type: DetectorTrainingRunner
+    key_metric: test.mAP@50:95
+    key_metric_asc: true
+    eval_interval: 10
+    optimizer:
+      type: AdamW
+      params:
+        lr: 0.0001
+    scheduler:
+      type: LRListScheduler
+      params:
+        milestones:
+        - 160
+        lr_list:
+        - - 1e-05
+    snapshots:
+      max_snapshots: 5
+      save_epochs: 25
+      save_optimizer_state: false
+  train_settings:
+    batch_size: 1
+    dataloader_workers: 0
+    dataloader_pin_memory: false
+    display_iters: 500
+    epochs: 250
+device: auto
+inference:
+  multithreading:
+    enabled: true
+    queue_length: 4
+    timeout: 30.0
+  compile:
+    enabled: false
+    backend: inductor
+  autocast:
+    enabled: false
+metadata:
+  project_path: ""
+  pose_config_path: ""
+  bodyparts:
+  - left_eye
+  - right_eye
+  - chin
+  - left_front_paw
+  - right_front_paw
+  - left_back_paw
+  - right_back_paw
+  - tail_base
+  - left_front_thigh
+  - right_front_thigh
+  - left_back_thigh
+  - right_back_thigh
+  - left_shoulder
+  - right_shoulder
+  - left_front_knee
+  - right_front_knee
+  - left_back_knee
+  - right_back_knee
+  - neck_base
+  - tail_mid
+  - left_ear_base
+  - right_ear_base
+  - left_mouth_corner
+  - right_mouth_corner
+  - nose
+  - tail_tip_first
+  unique_bodyparts: []
+  individuals:
+  - individual000
+  with_identity: false
+method: td
+model:
+  backbone:
+    type: HRNet
+    model_name: hrnet_w32
+    freeze_bn_stats: true
+    freeze_bn_weights: false
+    interpolate_branches: false
+    increased_channel_count: false
+  backbone_output_channels: 32
+  heads:
+    bodypart:
+      type: HeatmapHead
+      weight_init: normal
+      predictor:
+        type: HeatmapPredictor
+        apply_sigmoid: false
+        clip_scores: true
+        location_refinement: true
+        locref_std: 7.2801
+      target_generator:
+        type: HeatmapGaussianGenerator
+        num_heatmaps: 26
+        pos_dist_thresh: 17
+        heatmap_mode: KEYPOINT
+        gradient_masking: true
+        background_weight: 0.0
+        generate_locref: true
+        locref_std: 7.2801
+      criterion:
+        heatmap:
+          type: WeightedMSECriterion
+          weight: 1.0
+        locref:
+          type: WeightedHuberCriterion
+          weight: 0.05
+      heatmap_config:
+        channels:
+        - 32
+        kernel_size: []
+        strides: []
+        final_conv:
+          out_channels: 26
+          kernel_size: 1
+      locref_config:
+        channels:
+        - 32
+        kernel_size: []
+        strides: []
+        final_conv:
+          out_channels: 52
+          kernel_size: 1
+net_type: hrnet_w32
+runner:
+  type: PoseTrainingRunner
+  gpus:
+  key_metric: test.mAP
+  key_metric_asc: true
+  eval_interval: 10
+  optimizer:
+    type: AdamW
+    params:
+      lr: 0.0001
+  scheduler:
+    type: LRListScheduler
+    params:
+      lr_list:
+      - - 1e-05
+      - - 1e-06
+      milestones:
+      - 160
+      - 190
+  snapshots:
+    max_snapshots: 5
+    save_epochs: 10
+    save_optimizer_state: false
+train_settings:
+  batch_size: 64
+  dataloader_workers: 8
+  dataloader_pin_memory: false
+  display_iters: 500
+  epochs: 200
+  seed: 42
diff --git a/demo.py b/demo.py
index 2a68d82..bdf16db 100644
--- a/demo.py
+++ b/demo.py
@@ -21,7 +21,8 @@
 from prima.models import load_prima
 from prima.utils import recursive_to
 from prima.datasets.vitdet_dataset import ViTDetDataset, DEFAULT_MEAN, DEFAULT_STD
-from prima.utils.renderer import Renderer, cam_crop_to_full
+from prima.utils.detection import select_animal_boxes
+from prima.utils.weights import DEFAULT_HF_REPO_ID, resolve_prima_checkpoint_path
 import detectron2
 from detectron2 import model_zoo
 import warnings
@@ -29,17 +30,37 @@
 
 LIGHT_BLUE = (0.65098039, 0.74117647, 0.85882353)
 GREEN = (0.65, 0.86, 0.74)
+REPO_ROOT = Path(__file__).resolve().parent
 
 
+def load_renderer_components():
+    try:
+        from prima.utils.renderer import Renderer, cam_crop_to_full
+    except Exception as exc:
+        raise RuntimeError(
+            "Cannot initialize the PRIMA renderer. Rendering requires a working "
+            "pyrender/OpenGL backend such as EGL or OSMesa. Install the missing "
+            "OpenGL runtime for this environment, or run in an environment where "
+            "PYOPENGL_PLATFORM=egl/osmesa works."
+        ) from exc
+    return Renderer, cam_crop_to_full
+
 
 def main():
     parser = argparse.ArgumentParser(description='prima demo code')
-    parser.add_argument('--checkpoint', type=str,
-                        help='Path to pretrained model checkpoint')
+    parser.add_argument('--checkpoint', type=str, default='',
+                        help='Path to pretrained model checkpoint. Empty -> auto-download the default Stage 1 checkpoint.')
+    parser.add_argument('--hf-repo-id', '--hf_repo_id', dest='hf_repo_id',
+                        type=str, default=os.environ.get("PRIMA_HF_REPO_ID", DEFAULT_HF_REPO_ID),
+                        help='Hugging Face repo ID containing PRIMA demo assets')
+    parser.add_argument('--no-auto-download', '--no_auto_download', dest='no_auto_download', action='store_true',
+                        help='Disable automatic download of missing PRIMA demo assets')
     parser.add_argument('--img_folder', type=str, default='demo_data/', help='Folder with input images')
     parser.add_argument('--out_folder', type=str, default='demo_out', help='Output folder to save rendered results')
     parser.add_argument('--side_view', dest='side_view', action='store_true', default=False,
                         help='If set, render side view also')
+    parser.add_argument('--render_depth', dest='render_depth', action='store_true', default=False,
+                        help='If set, render depth map also')
     parser.add_argument('--save_mesh', dest='save_mesh', action='store_true', default=False,
                         help='If set, save meshes to disk also')
     parser.add_argument('--batch_size', type=int, default=1, help='Batch size for inference/fitting')
@@ -48,13 +69,21 @@ def main():
 
     args = parser.parse_args()
 
-    model, model_cfg = load_prima(args.checkpoint)
+    checkpoint_path = resolve_prima_checkpoint_path(
+        args.checkpoint,
+        data_dir=REPO_ROOT / "data",
+        auto_download=not args.no_auto_download,
+        hf_repo_id=args.hf_repo_id,
+    )
+
+    model, model_cfg = load_prima(checkpoint_path)
 
     device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
     model = model.to(device)
     model.eval()
 
     # Setup the renderer
+    Renderer, cam_crop_to_full = load_renderer_components()
     renderer = Renderer(model_cfg, faces=model.smal.faces)
 
     # Make output directory if it does not exist
@@ -63,22 +92,32 @@ def main():
     # Load detector
     cfg = detectron2.config.get_cfg()
     cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"))
-    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 
-    cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x/139173657/model_final_68b088.pkl" 
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
+    cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x/139173657/model_final_68b088.pkl"
+    cfg.MODEL.DEVICE = device.type
     detector = detectron2.engine.DefaultPredictor(cfg)
 
     img_paths = sorted([img for end in args.file_type for img in Path(args.img_folder).glob(end)])
+    num_readable_images = 0
+    num_rendered_results = 0
+    num_suppressed_detections = 0
     for img_path in img_paths:
         img_bgr = cv2.imread(str(img_path))
         if img_bgr is None:
             print(f"[WARN] Cannot read image: {img_path}")
             continue
+        num_readable_images += 1
         # Detect animals in image
         det_out = detector(img_bgr)
 
         det_instances = det_out['instances']
-        valid_idx = [i for i, (c, s) in enumerate(zip(det_instances.pred_classes, det_instances.scores)) if ((c in [15, 16, 17, 18, 19, 21, 22]) & (s > 0.7))]
-        boxes = det_instances.pred_boxes.tensor[valid_idx].cpu().numpy()
+        boxes, suppressed = select_animal_boxes(det_instances, score_threshold=0.7)
+        num_suppressed_detections += suppressed
+        if suppressed > 0:
+            print(f"[INFO] Suppressed {suppressed} duplicate animal detection(s) in {img_path}")
+        if len(boxes) == 0:
+            print(f"[INFO] No animal detected in {img_path}")
+            continue
 
         # Run PRIMA on detected animals
         dataset = ViTDetDataset(model_cfg, img_bgr, boxes)
@@ -125,9 +164,39 @@ def main():
                                         scene_bg_color=(1, 1, 1),
                                         side_view=True)
                     final_img = np.concatenate([final_img, side_img], axis=1)
+                    
+                if args.render_depth:
+                    depth_img = renderer(out['pred_vertices'][n].detach().cpu().numpy(),
+                                        out['pred_cam_t'][n].detach().cpu().numpy(),
+                                        white_img,
+                                        mesh_base_color=GREEN,
+                                        scene_bg_color=(1, 1, 1),
+                                        depth=True)
+
+                    valid_mask = depth_img > 0
+                    if np.sum(valid_mask) == 0:
+                        # no valid depth
+                        depth_norm = np.zeros_like(depth_img)
+                    else:
+                        min_val = np.min(depth_img[valid_mask])
+                        max_val = np.max(depth_img[valid_mask])
+                        if min_val == max_val:
+                            depth_norm = np.zeros_like(depth_img)
+                        else:
+                            depth_norm = (depth_img - min_val) / (max_val - min_val + 1e-8)
+                    depth_norm[~valid_mask] = 0
+                    
+                    depth_vis = (depth_norm * 255).astype(np.uint8)
+                    depth_vis = cv2.applyColorMap(depth_vis, cv2.COLORMAP_VIRIDIS)
+                    depth_vis = cv2.cvtColor(depth_vis, cv2.COLOR_BGR2RGB)
+                    depth_vis = depth_vis.astype(np.float32) / 255.0
+                    depth_vis[~valid_mask] = 0                    
+                    final_img = np.concatenate([final_img, depth_vis], axis=1)
+
 
                 cv2.imwrite(os.path.join(args.out_folder, f'{img_fn}_{animal_id}.png'), 
                             cv2.cvtColor((255 * final_img).astype(np.uint8), cv2.COLOR_RGB2BGR))
+                num_rendered_results += 1
 
                 # Add all verts and cams to list
                 verts = out['pred_vertices'][n].detach().cpu().numpy()
@@ -139,6 +208,13 @@ def main():
                     tmesh = renderer.vertices_to_trimesh(verts, camera_translation, LIGHT_BLUE)
                     tmesh.export(os.path.join(args.out_folder, f'{img_fn}_{animal_id}.obj'))
 
+    print(
+        f"[done] Demo complete. Processed {num_readable_images}/{len(img_paths)} image(s), "
+        f"saved {num_rendered_results} rendered result(s) to {args.out_folder}."
+    )
+    if num_suppressed_detections > 0:
+        print(f"[done] Suppressed {num_suppressed_detections} duplicate animal detection(s).")
+
 
 if __name__ == '__main__':
     main()
diff --git a/demo.sh b/demo.sh
new file mode 100644
index 0000000..6dace89
--- /dev/null
+++ b/demo.sh
@@ -0,0 +1,12 @@
+# Default PRIMA Stage 1 inference checkpoint:
+#   data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt
+#
+# If this local file is missing, it will be downloaded from the PRIMA Hugging Face repo.
+# To use another local checkpoint instead, update this path.
+# For example: checkpoint='data/PRIMAS3/checkpoints/s3ckpt_inference.ckpt'
+checkpoint='data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt'
+
+python demo.py \
+  --checkpoint "${checkpoint}" \
+  --img_folder demo_data/ \
+  --out_folder demo_out/
diff --git a/demo_data/000000015956_horse.png b/demo_data/000000015956_horse.png
index 2b05a48..15fab77 100644
Binary files a/demo_data/000000015956_horse.png and b/demo_data/000000015956_horse.png differ
diff --git a/demo_data/000000315905_zebra.jpg b/demo_data/000000315905_zebra.jpg
index 857dc57..c476bdf 100644
Binary files a/demo_data/000000315905_zebra.jpg and b/demo_data/000000315905_zebra.jpg differ
diff --git a/demo_data/beagle.jpg b/demo_data/beagle.jpg
index 4fd9747..e3ee922 100644
Binary files a/demo_data/beagle.jpg and b/demo_data/beagle.jpg differ
diff --git a/demo_data/n02101388_1188.png b/demo_data/n02101388_1188.png
index b424bab..873d586 100644
Binary files a/demo_data/n02101388_1188.png and b/demo_data/n02101388_1188.png differ
diff --git a/demo_data/n02412080_12159.png b/demo_data/n02412080_12159.png
index 36e700f..a6517fa 100644
Binary files a/demo_data/n02412080_12159.png and b/demo_data/n02412080_12159.png differ
diff --git a/demo_data/shepherd_hati.jpg b/demo_data/shepherd_hati.jpg
index 7e770f2..4fd79ba 100644
Binary files a/demo_data/shepherd_hati.jpg and b/demo_data/shepherd_hati.jpg differ
diff --git a/demo_tta.py b/demo_tta.py
index 3980d3d..7c713bf 100644
--- a/demo_tta.py
+++ b/demo_tta.py
@@ -8,24 +8,20 @@
 """
 
 """
-demo_tta.py: PRIMA inference with DeepLabCut SuperAnimal TTA
+demo_tta.py: PRIMA inference with fine-tuned DeepLabCut SuperAnimal TTA
 
 Pipeline:
 1. Run Detectron2 to detect animals in the input image.
 2. Run PRIMA on each detected animal to obtain 3D pose/shape estimation.
-3. Run DeepLabCut SuperAnimal to obtain 2D keypoint estimation.
-4. Map the 39 SuperAnimal keypoints to the 26 PRIMA keypoints.
-5. Run test-time adaptation (TTA) with user-specified lr and num_iters
+3. Run a fine-tuned DeepLabCut SuperAnimal pose model (Animal3D 26-joint
+   layout) to obtain 2D keypoints already in PRIMA topology. The fine-tuned
+   snapshot is wired into DLC's
+   ``superanimal_analyze_images`` via the ``customized_pose_checkpoint``
+   and ``customized_model_config`` kwargs.
+4. Run test-time adaptation (TTA) with user-specified lr and num_iters
    to further optimize the 3D pose and shape estimation.
-6. Render and save before/after TTA results (PNG + OBJ) and the
+5. Render and save before/after TTA results (PNG + OBJ) and the
    26-keypoint visualization (PNG).
-
-Reference code:
-- Test-time adaptation: prima/../eval_with_tta.py
-- DeepLabCut: https://github.com/AdaptiveMotorControlLab/FMPose3D/blob/main/animals/demo/vis_animals.py
-- Keypoint mapping (SuperAnimal 39 → PRIMA 26):
-    keypoint_mapping = {"quadruped80k":[10, 5, -1, 26, 29, 30, 35, 22, 24, 27, 31, 32, -1, -1,
-                                     25, 28, 33, 34, 15, 23, 11, 6, 4, 3, 0, -1]}
 """
 
 
@@ -43,25 +39,31 @@
 import torch.utils.data
 from tqdm import tqdm
 
-import detectron2
-import detectron2.config
-import detectron2.engine
-from detectron2 import model_zoo
-
 from prima.models import load_prima
 from prima.utils import recursive_to
 from prima.datasets.vitdet_dataset import ViTDetDataset, DEFAULT_MEAN, DEFAULT_STD
-from prima.utils.renderer import Renderer, cam_crop_to_full
+from prima.utils.detection import ANIMAL_COCO_IDS, select_animal_boxes
+from prima.utils.weights import DEFAULT_HF_REPO_ID, resolve_prima_checkpoint_path
 
 warnings.filterwarnings("ignore")
 
 LIGHT_BLUE = (0.65098039, 0.74117647, 0.85882353)
 GREEN = (0.65, 0.86, 0.74)
 
-ANIMAL_COCO_IDS = [15, 16, 17, 18, 19, 21, 22]
-keypoint_mapping = {
-    "quadruped80k": [10, 5, -1, 26, 29, 30, 35, 22, 24, 27, 31, 32, -1, -1, 25, 28, 33, 34, 15, 23, 11, 6, 4, 3, 0, -1]
-}
+REPO_ROOT = Path(__file__).resolve().parent
+
+
+def load_renderer_components():
+    try:
+        from prima.utils.renderer import Renderer, cam_crop_to_full
+    except Exception as exc:
+        raise RuntimeError(
+            "Cannot initialize the PRIMA renderer. Rendering requires a working "
+            "pyrender/OpenGL backend such as EGL or OSMesa. Install the missing "
+            "OpenGL runtime for this environment, or run in an environment where "
+            "PYOPENGL_PLATFORM=egl/osmesa works."
+        ) from exc
+    return Renderer, cam_crop_to_full
 
 
 def denorm_patch_to_rgb(img_tensor: torch.Tensor) -> np.ndarray:
@@ -70,18 +72,6 @@ def denorm_patch_to_rgb(img_tensor: torch.Tensor) -> np.ndarray:
     return np.clip(patch, 0.0, 1.0)
 
 
-def map_superanimal_to_prima(bodyparts_xyc: np.ndarray) -> np.ndarray:
-    mapping = keypoint_mapping["quadruped80k"]
-    num_src = bodyparts_xyc.shape[0]
-    mapped = np.zeros((len(mapping), 3), dtype=np.float32)
-
-    for tgt_i, src_i in enumerate(mapping):
-        if src_i >= 0 and src_i < num_src:
-            mapped[tgt_i] = bodyparts_xyc[src_i]
-
-    return mapped
-
-
 def save_keypoint_vis(patch_rgb: np.ndarray, kpts_xyc: np.ndarray, save_path: str) -> None:
     vis = cv2.cvtColor((patch_rgb * 255).astype(np.uint8), cv2.COLOR_RGB2BGR).copy()
     num_kpts = len(kpts_xyc)
@@ -102,7 +92,63 @@ def save_keypoint_vis(patch_rgb: np.ndarray, kpts_xyc: np.ndarray, save_path: st
     cv2.imwrite(save_path, vis)
 
 
+def depth_to_viridis_rgb(depth_img: np.ndarray) -> np.ndarray:
+    valid_mask = depth_img > 0
+    if np.sum(valid_mask) == 0:
+        depth_norm = np.zeros_like(depth_img)
+    else:
+        min_val = np.min(depth_img[valid_mask])
+        max_val = np.max(depth_img[valid_mask])
+        if min_val == max_val:
+            depth_norm = np.zeros_like(depth_img)
+        else:
+            depth_norm = (depth_img - min_val) / (max_val - min_val + 1e-8)
+    depth_norm[~valid_mask] = 0
+
+    depth_vis = (depth_norm * 255).astype(np.uint8)
+    depth_vis = cv2.applyColorMap(depth_vis, cv2.COLORMAP_VIRIDIS)
+    depth_vis = cv2.cvtColor(depth_vis, cv2.COLOR_BGR2RGB)
+    depth_vis = depth_vis.astype(np.float32) / 255.0
+    depth_vis[~valid_mask] = 0
+    return depth_vis
+
+
+def resolve_sa_weights_path(local_path: str) -> str:
+    """Return a local path to the fine-tuned SuperAnimal .pt snapshot.
+
+    If ``local_path`` is empty, downloads ``sa_finetune_hrnet_w32.pt`` from the
+    ``MLAdaptiveIntelligence/FMPose3D`` Hugging Face repo (cached under
+    ``~/.cache/huggingface``).
+    """
+    if local_path:
+        return local_path
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError:
+        raise ImportError(
+            "huggingface_hub is required to auto-download the fine-tuned "
+            "SuperAnimal weights. Install with `pip install huggingface_hub`, "
+            "or pass --saved_2d_model_path with a local .pt file."
+        ) from None
+    repo_id = "MLAdaptiveIntelligence/FMPose3D"
+    filename = "sa_finetune_hrnet_w32.pt"
+    try:
+        cached_path = hf_hub_download(repo_id=repo_id, filename=filename, local_files_only=True)
+    except Exception:
+        print(f"No --saved_2d_model_path provided; downloading '{filename}' from {repo_id}...")
+        return hf_hub_download(repo_id=repo_id, filename=filename)
+
+    print(f"Using cached SuperAnimal weights: {cached_path}")
+    return cached_path
+
+
 def run_superanimal_on_patch(patch_rgb: np.ndarray, args, tmp_dir: str):
+    """Predict 26-joint 2D keypoints on a single PRIMA patch using a
+    fine-tuned DeepLabCut SuperAnimal snapshot.
+
+    Returns an ``(26, 3)`` array of ``(x, y, confidence)`` in patch
+    pixel coordinates, or ``None`` if no individual was detected.
+    """
     try:
         from deeplabcut.pose_estimation_pytorch.apis import superanimal_analyze_images
     except Exception as e:
@@ -113,13 +159,17 @@ def run_superanimal_on_patch(patch_rgb: np.ndarray, args, tmp_dir: str):
     patch_path = os.path.join(tmp_dir, "patch.png")
     cv2.imwrite(patch_path, cv2.cvtColor((patch_rgb * 255).astype(np.uint8), cv2.COLOR_RGB2BGR))
 
+    dlc_device = "cuda" if torch.cuda.is_available() else "cpu"
     preds = superanimal_analyze_images(
-        args.superanimal_name,
-        args.superanimal_model_name,
-        args.superanimal_detector_name,
-        patch_path,
-        args.superanimal_max_individuals,
+        superanimal_name=args.superanimal_name,
+        model_name=args.superanimal_model_name,
+        detector_name=args.superanimal_detector_name,
+        images=patch_path,
+        max_individuals=args.superanimal_max_individuals,
         out_folder=tmp_dir,
+        device=dlc_device,
+        customized_model_config=args.pytorch_config_2d_path,
+        customized_pose_checkpoint=args.saved_2d_model_path,
     )
 
     payload = preds.get(patch_path, None)
@@ -130,16 +180,28 @@ def run_superanimal_on_patch(patch_rgb: np.ndarray, args, tmp_dir: str):
         return None
 
     best_idx = int(np.argmax(bodyparts[..., 2].mean(axis=1)))
-    return bodyparts[best_idx]
-
-
-def render_and_save(renderer, out, batch, img_fn, animal_id, out_folder, suffix, side_view, save_mesh):
+    return bodyparts[best_idx].astype(np.float32)
+
+
+def render_and_save(
+    renderer,
+    cam_crop_to_full_fn,
+    out,
+    batch,
+    img_fn,
+    animal_id,
+    out_folder,
+    suffix,
+    side_view,
+    save_mesh,
+    render_depth=False,
+):
     pred_cam = out['pred_cam']
     box_center = batch['box_center'].float()
     box_size = batch['box_size'].float()
     img_size = batch['img_size'].float()
     scaled_focal_length = batch['focal_length'][0, 0] / batch['img'].shape[-1] * img_size.max()
-    pred_cam_t_full = cam_crop_to_full(pred_cam, box_center, box_size, img_size, scaled_focal_length)
+    pred_cam_t_full = cam_crop_to_full_fn(pred_cam, box_center, box_size, img_size, scaled_focal_length)
 
     white_img = (torch.ones_like(batch['img'][0]).cpu() - DEFAULT_MEAN[:, None, None] / 255) / (
         DEFAULT_STD[:, None, None] / 255
@@ -166,6 +228,17 @@ def render_and_save(renderer, out, batch, img_fn, animal_id, out_folder, suffix,
         )
         final_img = np.concatenate([final_img, side_img], axis=1)
 
+    if render_depth:
+        depth_img = renderer(
+            out['pred_vertices'][0].detach().cpu().numpy(),
+            out['pred_cam_t'][0].detach().cpu().numpy(),
+            white_img,
+            mesh_base_color=GREEN,
+            scene_bg_color=(1, 1, 1),
+            depth=True,
+        )
+        final_img = np.concatenate([final_img, depth_to_viridis_rgb(depth_img)], axis=1)
+
     cv2.imwrite(
         os.path.join(out_folder, f'{img_fn}_{animal_id}_{suffix}.png'),
         cv2.cvtColor((255 * final_img).astype(np.uint8), cv2.COLOR_RGB2BGR),
@@ -212,11 +285,18 @@ def tta_optimize(model, batch, gt_kpts_norm, num_iters, lr):
 
 def main():
     parser = argparse.ArgumentParser(description='PRIMA + SuperAnimal + TTA demo')
-    parser.add_argument('--checkpoint', type=str, required=True, help='Path to pretrained PRIMA checkpoint')
+    parser.add_argument('--checkpoint', type=str, default='',
+                        help='Path to pretrained PRIMA checkpoint. Empty -> auto-download the default Stage 1 checkpoint.')
+    parser.add_argument('--hf-repo-id', '--hf_repo_id', dest='hf_repo_id',
+                        type=str, default=os.environ.get("PRIMA_HF_REPO_ID", DEFAULT_HF_REPO_ID),
+                        help='Hugging Face repo ID containing PRIMA demo assets')
+    parser.add_argument('--no-auto-download', '--no_auto_download', dest='no_auto_download', action='store_true',
+                        help='Disable automatic download of missing PRIMA demo assets')
     parser.add_argument('--img_path', type=str, default=None, help='Single image path')
     parser.add_argument('--img_folder', type=str, default='demo_data/', help='Folder with input images')
     parser.add_argument('--out_folder', type=str, default='demo_out_tta', help='Output folder')
     parser.add_argument('--side_view', dest='side_view', action='store_true', default=False, help='Render side view')
+    parser.add_argument('--render_depth', dest='render_depth', action='store_true', default=False, help='Render depth map')
     parser.add_argument('--save_mesh', dest='save_mesh', action='store_true', default=False, help='Save meshes')
     parser.add_argument('--file_type', nargs='+', default=['*.jpg', '*.png', '*.jpeg', '*.JPEG'], help='Image globs')
     parser.add_argument('--det_thresh', type=float, default=0.7, help='Detectron2 score threshold for animals')
@@ -229,21 +309,42 @@ def main():
     parser.add_argument('--superanimal_model_name', type=str, default='hrnet_w32')
     parser.add_argument('--superanimal_detector_name', type=str, default='fasterrcnn_resnet50_fpn_v2')
     parser.add_argument('--superanimal_max_individuals', type=int, default=1)
+    parser.add_argument('--saved_2d_model_path', type=str, default='',
+                        help='Path to the fine-tuned SuperAnimal 26-joint .pt snapshot. '
+                             'Empty -> auto-download sa_finetune_hrnet_w32.pt from '
+                             'MLAdaptiveIntelligence/FMPose3D on Hugging Face Hub.')
+    parser.add_argument('--pytorch_config_2d_path', type=str,
+                        default=str(Path(__file__).resolve().parent / 'configs' / 'sa_finetune_hrnet_w32.yaml'),
+                        help='Path to the DLC pytorch config yaml for the fine-tuned snapshot. '
+                             'Defaults to the bundled configs/sa_finetune_hrnet_w32.yaml.')
 
     args = parser.parse_args()
+    checkpoint_path = resolve_prima_checkpoint_path(
+        args.checkpoint,
+        data_dir=REPO_ROOT / "data",
+        auto_download=not args.no_auto_download,
+        hf_repo_id=args.hf_repo_id,
+    )
+    args.saved_2d_model_path = resolve_sa_weights_path(args.saved_2d_model_path)
 
-    model, model_cfg = load_prima(args.checkpoint)
+    model, model_cfg = load_prima(checkpoint_path)
     device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
     model = model.to(device)
     model.eval()
 
+    Renderer, cam_crop_to_full_fn = load_renderer_components()
     renderer = Renderer(model_cfg, faces=model.smal.faces)
     os.makedirs(args.out_folder, exist_ok=True)
 
+    import detectron2.config
+    import detectron2.engine
+    from detectron2 import model_zoo
+
     cfg = detectron2.config.get_cfg()
     cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"))
     cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
     cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x/139173657/model_final_68b088.pkl"
+    cfg.MODEL.DEVICE = device.type
     detector = detectron2.engine.DefaultPredictor(cfg)
 
     if args.img_path is not None:
@@ -258,11 +359,13 @@ def main():
             continue
         det_out = detector(img_bgr)
         det_instances = det_out['instances']
-        valid_idx = [
-            i for i, (c, s) in enumerate(zip(det_instances.pred_classes, det_instances.scores))
-            if (int(c) in ANIMAL_COCO_IDS) and (float(s) > args.det_thresh)
-        ]
-        boxes = det_instances.pred_boxes.tensor[valid_idx].cpu().numpy()
+        boxes, suppressed = select_animal_boxes(
+            det_instances,
+            animal_class_ids=ANIMAL_COCO_IDS,
+            score_threshold=args.det_thresh,
+        )
+        if suppressed > 0:
+            print(f"[INFO] Suppressed {suppressed} duplicate animal detection(s) in {img_path}")
 
         if len(boxes) == 0:
             print(f"[INFO] No animal detected in {img_path}")
@@ -281,6 +384,7 @@ def main():
 
             render_and_save(
                 renderer,
+                cam_crop_to_full_fn,
                 out_before,
                 batch,
                 img_fn,
@@ -289,31 +393,31 @@ def main():
                 suffix='before_tta',
                 side_view=args.side_view,
                 save_mesh=args.save_mesh,
+                render_depth=args.render_depth,
             )
 
             patch_rgb = denorm_patch_to_rgb(batch['img'][0])
             with tempfile.TemporaryDirectory(prefix=f"dlc_{img_fn}_{animal_id}_") as tmp_dir:
-                bodyparts_xyc = run_superanimal_on_patch(patch_rgb, args, tmp_dir)
+                kpts_xyc = run_superanimal_on_patch(patch_rgb, args, tmp_dir)
 
-            if bodyparts_xyc is None:
+            if kpts_xyc is None:
                 print(f"[WARN] No SuperAnimal keypoints for {img_fn}_{animal_id}, skip TTA")
                 continue
 
-            mapped_xyc = map_superanimal_to_prima(bodyparts_xyc)
-            mapped_xyc[mapped_xyc[:, 2] < args.kp_conf_thresh, 2] = 0.0
+            kpts_xyc[kpts_xyc[:, 2] < args.kp_conf_thresh, 2] = 0.0
 
             save_keypoint_vis(
                 patch_rgb,
-                mapped_xyc,
+                kpts_xyc,
                 os.path.join(args.out_folder, f"{img_fn}_{animal_id}_prima26_kpts.png"),
             )
-            np.save(os.path.join(args.out_folder, f"{img_fn}_{animal_id}_prima26_kpts.npy"), mapped_xyc)
+            np.save(os.path.join(args.out_folder, f"{img_fn}_{animal_id}_prima26_kpts.npy"), kpts_xyc)
 
             patch_h, patch_w = patch_rgb.shape[:2]
-            mapped_norm = mapped_xyc.copy()
-            mapped_norm[:, 0] = mapped_norm[:, 0] / float(patch_w) - 0.5
-            mapped_norm[:, 1] = mapped_norm[:, 1] / float(patch_h) - 0.5
-            gt_kpts_norm = torch.from_numpy(mapped_norm[None]).to(device=device, dtype=batch['img'].dtype)
+            kpts_norm = kpts_xyc.copy()
+            kpts_norm[:, 0] = kpts_norm[:, 0] / float(patch_w) - 0.5
+            kpts_norm[:, 1] = kpts_norm[:, 1] / float(patch_h) - 0.5
+            gt_kpts_norm = torch.from_numpy(kpts_norm[None]).to(device=device, dtype=batch['img'].dtype)
 
             out_after = tta_optimize(
                 model,
@@ -325,6 +429,7 @@ def main():
 
             render_and_save(
                 renderer,
+                cam_crop_to_full_fn,
                 out_after,
                 batch,
                 img_fn,
@@ -333,9 +438,9 @@ def main():
                 suffix='after_tta',
                 side_view=args.side_view,
                 save_mesh=args.save_mesh,
+                render_depth=args.render_depth,
             )
 
 
 if __name__ == '__main__':
     main()
-
diff --git a/demo_tta.sh b/demo_tta.sh
new file mode 100644
index 0000000..6249a2d
--- /dev/null
+++ b/demo_tta.sh
@@ -0,0 +1,16 @@
+
+# Empty checkpoint uses the default PRIMA Stage 1 inference checkpoint:
+#   data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt
+#
+# This standard path is auto-downloaded from the PRIMA Hugging Face repo if missing.
+# To use another local checkpoint instead, update this path.
+# For example: checkpoint='data/PRIMAS3/checkpoints/s3ckpt_inference.ckpt'
+checkpoint='data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt'
+
+python3 demo_tta.py \
+  --checkpoint "${checkpoint}" \
+  --img_folder demo_data/ \
+  --out_folder demo_out_tta/ \
+  --tta_lr 1e-6 \
+  --tta_num_iters 30 \
+  --render_depth
diff --git a/demo_video.py b/demo_video.py
new file mode 100644
index 0000000..fa97bdd
--- /dev/null
+++ b/demo_video.py
@@ -0,0 +1,353 @@
+"""
+PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation
+
+Official implementation of the paper:
+"PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation"
+by Xiaohang Yu, Ti Wang, and Mackenzie Weygandt Mathis
+Licensed under a modified MIT license
+"""
+
+from pathlib import Path
+import argparse
+import os
+import warnings
+
+import cv2
+import detectron2
+import detectron2.config
+import detectron2.engine
+import numpy as np
+import torch
+import torch.utils
+import torch.utils.data
+from detectron2 import model_zoo
+from tqdm import tqdm
+
+from prima.datasets.vitdet_dataset import ViTDetDataset, DEFAULT_MEAN, DEFAULT_STD
+from prima.models import load_prima
+from prima.utils import recursive_to
+from prima.utils.detection import ANIMAL_COCO_IDS
+from prima.utils.weights import DEFAULT_HF_REPO_ID, resolve_prima_checkpoint_path
+
+warnings.filterwarnings("ignore")
+
+LIGHT_BLUE = (0.65098039, 0.74117647, 0.85882353)
+GREEN = (0.65, 0.86, 0.74)
+REPO_ROOT = Path(__file__).resolve().parent
+
+
+def load_renderer_components():
+    try:
+        from prima.utils.renderer import Renderer, cam_crop_to_full
+    except Exception as exc:
+        raise RuntimeError(
+            "Cannot initialize the PRIMA renderer. Rendering requires a working "
+            "pyrender/OpenGL backend such as EGL or OSMesa. Install the missing "
+            "OpenGL runtime for this environment, or run in an environment where "
+            "PYOPENGL_PLATFORM=egl/osmesa works."
+        ) from exc
+    return Renderer, cam_crop_to_full
+
+
+def select_top_confidence_animal_box(det_instances, score_threshold=0.7):
+    classes = det_instances.pred_classes.detach().cpu().numpy()
+    scores = det_instances.scores.detach().cpu().numpy()
+    class_ids = set(int(class_id) for class_id in ANIMAL_COCO_IDS)
+    valid_idx = np.array(
+        [
+            i
+            for i, (class_id, score) in enumerate(zip(classes, scores))
+            if int(class_id) in class_ids and float(score) > float(score_threshold)
+        ],
+        dtype=np.int64,
+    )
+    if len(valid_idx) == 0:
+        return np.zeros((0, 4), dtype=np.float32), None
+
+    top_idx = valid_idx[int(np.argmax(scores[valid_idx]))]
+    box = det_instances.pred_boxes.tensor[top_idx].detach().cpu().numpy().astype(np.float32)
+    return box[None], float(scores[top_idx])
+
+
+def depth_to_viridis_rgb(depth_img):
+    valid_mask = depth_img > 0
+    if np.sum(valid_mask) == 0:
+        depth_norm = np.zeros_like(depth_img)
+    else:
+        min_val = np.min(depth_img[valid_mask])
+        max_val = np.max(depth_img[valid_mask])
+        if min_val == max_val:
+            depth_norm = np.zeros_like(depth_img)
+        else:
+            depth_norm = (depth_img - min_val) / (max_val - min_val + 1e-8)
+    depth_norm[~valid_mask] = 0
+
+    depth_vis = (depth_norm * 255).astype(np.uint8)
+    depth_vis = cv2.applyColorMap(depth_vis, cv2.COLORMAP_VIRIDIS)
+    depth_vis = cv2.cvtColor(depth_vis, cv2.COLOR_BGR2RGB)
+    depth_vis = depth_vis.astype(np.float32) / 255.0
+    depth_vis[~valid_mask] = 0
+    return depth_vis
+
+
+def make_empty_output_frame(frame_bgr, img_res, num_panels):
+    frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+    first_panel = cv2.resize(frame_rgb, (img_res, img_res)).astype(np.float32) / 255.0
+    blank_panel = np.ones((img_res, img_res, 3), dtype=np.float32)
+    panels = [first_panel] + [blank_panel.copy() for _ in range(num_panels - 1)]
+    return np.concatenate(panels, axis=1)
+
+
+def make_full_frame_output(frame_bgr):
+    return cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
+
+
+def get_video_rotation(cap, rotate_arg):
+    if rotate_arg != "auto":
+        return rotate_arg
+    orientation_prop = getattr(cv2, "CAP_PROP_ORIENTATION_META", None)
+    if orientation_prop is None:
+        return "none"
+    orientation = int(cap.get(orientation_prop) or 0) % 360
+    if orientation == 90:
+        return "90cw"
+    if orientation == 180:
+        return "180"
+    if orientation == 270:
+        return "90ccw"
+    return "none"
+
+
+def rotate_frame(frame_bgr, rotation):
+    if rotation == "90cw":
+        return cv2.rotate(frame_bgr, cv2.ROTATE_90_CLOCKWISE)
+    if rotation == "90ccw":
+        return cv2.rotate(frame_bgr, cv2.ROTATE_90_COUNTERCLOCKWISE)
+    if rotation == "180":
+        return cv2.rotate(frame_bgr, cv2.ROTATE_180)
+    return frame_bgr
+
+
+def main():
+    parser = argparse.ArgumentParser(description="PRIMA video demo")
+    parser.add_argument("--checkpoint", type=str, default="",
+                        help="Path to pretrained model checkpoint. Empty -> auto-download the default Stage 1 checkpoint.")
+    parser.add_argument("--hf-repo-id", "--hf_repo_id", dest="hf_repo_id",
+                        type=str, default=os.environ.get("PRIMA_HF_REPO_ID", DEFAULT_HF_REPO_ID),
+                        help="Hugging Face repo ID containing PRIMA demo assets")
+    parser.add_argument("--no-auto-download", "--no_auto_download", dest="no_auto_download", action="store_true",
+                        help="Disable automatic download of missing PRIMA demo assets")
+    parser.add_argument("--video_path", type=str, required=True, help="Input video path")
+    parser.add_argument("--out_video", type=str, default="demo_video_out.mp4", help="Output rendered video path")
+    parser.add_argument("--out_folder", type=str, default="demo_video_out", help="Output folder for optional meshes")
+    parser.add_argument("--det_thresh", type=float, default=0.7, help="Animal detection confidence threshold")
+    parser.add_argument("--side_view", dest="side_view", action="store_true", default=False,
+                        help="If set, render side view also")
+    parser.add_argument("--render_depth", dest="render_depth", action="store_true", default=False,
+                        help="If set, render depth map also")
+    parser.add_argument("--full_frame", dest="full_frame", action="store_true", default=False,
+                        help="Render the mesh overlay on the full video frame instead of crop-panel output")
+    parser.add_argument("--save_mesh", dest="save_mesh", action="store_true", default=False,
+                        help="If set, save one mesh per processed frame")
+    parser.add_argument("--max_frames", type=int, default=-1,
+                        help="Maximum number of frames to process. Use -1 for the full video.")
+    parser.add_argument("--frame_stride", type=int, default=1,
+                        help="Process every Nth frame. Output video contains processed frames only.")
+    parser.add_argument("--rotate", type=str, default="auto",
+                        choices=["auto", "none", "90cw", "90ccw", "180"],
+                        help="Rotate input frames before detection/rendering. "
+                             "auto uses video orientation metadata when OpenCV exposes it.")
+
+    args = parser.parse_args()
+    os.makedirs(args.out_folder, exist_ok=True)
+    out_video_parent = os.path.dirname(args.out_video)
+    if out_video_parent:
+        os.makedirs(out_video_parent, exist_ok=True)
+
+    checkpoint_path = resolve_prima_checkpoint_path(
+        args.checkpoint,
+        data_dir=REPO_ROOT / "data",
+        auto_download=not args.no_auto_download,
+        hf_repo_id=args.hf_repo_id,
+    )
+
+    model, model_cfg = load_prima(checkpoint_path)
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    model = model.to(device)
+    model.eval()
+
+    Renderer, cam_crop_to_full = load_renderer_components()
+    renderer = Renderer(model_cfg, faces=model.smal.faces)
+
+    cfg = detectron2.config.get_cfg()
+    cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"))
+    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
+    cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x/139173657/model_final_68b088.pkl"
+    cfg.MODEL.DEVICE = device.type
+    detector = detectron2.engine.DefaultPredictor(cfg)
+
+    cap = cv2.VideoCapture(args.video_path)
+    if not cap.isOpened():
+        raise RuntimeError(f"Cannot open video: {args.video_path}")
+    orientation_auto_prop = getattr(cv2, "CAP_PROP_ORIENTATION_AUTO", None)
+    if orientation_auto_prop is not None:
+        cap.set(orientation_auto_prop, 0)
+    frame_rotation = get_video_rotation(cap, args.rotate)
+    print(f"[video] frame rotation: {frame_rotation}")
+
+    src_fps = cap.get(cv2.CAP_PROP_FPS)
+    fps = src_fps / max(1, args.frame_stride) if src_fps and src_fps > 0 else 30.0
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
+    if args.max_frames > 0:
+        total_steps = min(total_frames, args.max_frames * max(1, args.frame_stride))
+    else:
+        total_steps = total_frames
+
+    img_res = int(model_cfg.MODEL.IMAGE_SIZE)
+    num_panels = 2 + int(args.side_view) + int(args.render_depth)
+    writer = None
+    out_size = None
+
+    video_stem = Path(args.video_path).stem
+    frame_idx = 0
+    processed_frames = 0
+    rendered_frames = 0
+    skipped_frames = 0
+
+    pbar = tqdm(total=total_steps if total_steps > 0 else None, desc="Processing video")
+    try:
+        while True:
+            ret, frame_bgr = cap.read()
+            if not ret:
+                break
+            frame_bgr = rotate_frame(frame_bgr, frame_rotation)
+
+            if args.max_frames > 0 and processed_frames >= args.max_frames:
+                break
+
+            should_process = frame_idx % max(1, args.frame_stride) == 0
+            if not should_process:
+                frame_idx += 1
+                pbar.update(1)
+                continue
+
+            det_out = detector(frame_bgr)
+            boxes, top_score = select_top_confidence_animal_box(
+                det_out["instances"],
+                score_threshold=args.det_thresh,
+            )
+
+            if len(boxes) == 0:
+                if args.full_frame:
+                    final_img = make_full_frame_output(frame_bgr)
+                else:
+                    final_img = make_empty_output_frame(frame_bgr, img_res, num_panels)
+                skipped_frames += 1
+            else:
+                dataset = ViTDetDataset(model_cfg, frame_bgr, boxes)
+                batch = next(iter(torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)))
+                batch = recursive_to(batch, device)
+                with torch.no_grad():
+                    out = model(batch)
+
+                pred_cam = out["pred_cam"]
+                box_center = batch["box_center"].float()
+                box_size = batch["box_size"].float()
+                img_size = batch["img_size"].float()
+                scaled_focal_length = model_cfg.EXTRA.FOCAL_LENGTH / model_cfg.MODEL.IMAGE_SIZE * img_size.max()
+                pred_cam_t_full = cam_crop_to_full(
+                    pred_cam,
+                    box_center,
+                    box_size,
+                    img_size,
+                    scaled_focal_length,
+                ).detach().cpu().numpy()
+
+                if args.full_frame:
+                    frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+                    final_img = renderer(
+                        out["pred_vertices"][0].detach().cpu().numpy(),
+                        pred_cam_t_full[0],
+                        frame_rgb,
+                        full_frame=True,
+                        mesh_base_color=GREEN,
+                        scene_bg_color=(1, 1, 1),
+                        focal_length=float(scaled_focal_length.detach().cpu().numpy()),
+                    )
+                else:
+                    white_img = (torch.ones_like(batch["img"][0]).cpu() - DEFAULT_MEAN[:, None, None] / 255) / (
+                        DEFAULT_STD[:, None, None] / 255
+                    )
+                    input_patch = (
+                        batch["img"][0].cpu() * DEFAULT_STD[:, None, None] + DEFAULT_MEAN[:, None, None]
+                    ) / 255.0
+                    input_patch = input_patch.permute(1, 2, 0).numpy()
+
+                    regression_img = renderer(
+                        out["pred_vertices"][0].detach().cpu().numpy(),
+                        out["pred_cam_t"][0].detach().cpu().numpy(),
+                        batch["img"][0],
+                        mesh_base_color=GREEN,
+                        scene_bg_color=(1, 1, 1),
+                    )
+                    final_img = np.concatenate([input_patch, regression_img], axis=1)
+
+                    if args.side_view:
+                        side_img = renderer(
+                            out["pred_vertices"][0].detach().cpu().numpy(),
+                            out["pred_cam_t"][0].detach().cpu().numpy(),
+                            white_img,
+                            mesh_base_color=GREEN,
+                            scene_bg_color=(1, 1, 1),
+                            side_view=True,
+                        )
+                        final_img = np.concatenate([final_img, side_img], axis=1)
+
+                    if args.render_depth:
+                        depth_img = renderer(
+                            out["pred_vertices"][0].detach().cpu().numpy(),
+                            out["pred_cam_t"][0].detach().cpu().numpy(),
+                            white_img,
+                            mesh_base_color=GREEN,
+                            scene_bg_color=(1, 1, 1),
+                            depth=True,
+                        )
+                        final_img = np.concatenate([final_img, depth_to_viridis_rgb(depth_img)], axis=1)
+
+                if args.save_mesh:
+                    verts = out["pred_vertices"][0].detach().cpu().numpy()
+                    cam_t = pred_cam_t_full[0]
+                    tmesh = renderer.vertices_to_trimesh(verts, cam_t.copy(), LIGHT_BLUE)
+                    mesh_name = f"{video_stem}_frame{frame_idx:06d}_score{top_score:.3f}.obj"
+                    tmesh.export(os.path.join(args.out_folder, mesh_name))
+
+                rendered_frames += 1
+
+            frame_out = cv2.cvtColor((255 * final_img).astype(np.uint8), cv2.COLOR_RGB2BGR)
+            if writer is None:
+                out_size = (frame_out.shape[1], frame_out.shape[0])
+                writer = cv2.VideoWriter(args.out_video, cv2.VideoWriter_fourcc(*"mp4v"), fps, out_size)
+                if not writer.isOpened():
+                    raise RuntimeError(f"Cannot open output video writer: {args.out_video}")
+            elif (frame_out.shape[1], frame_out.shape[0]) != out_size:
+                frame_out = cv2.resize(frame_out, out_size)
+            writer.write(frame_out)
+
+            processed_frames += 1
+            frame_idx += 1
+            pbar.update(1)
+    finally:
+        pbar.close()
+        cap.release()
+        if writer is not None:
+            writer.release()
+
+    print(
+        f"[done] Processed {processed_frames} frame(s) from {args.video_path}; "
+        f"rendered {rendered_frames}, no-detection placeholders {skipped_frames}. "
+        f"Saved video to {args.out_video}."
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/demo_video.sh b/demo_video.sh
new file mode 100644
index 0000000..b9f8a6e
--- /dev/null
+++ b/demo_video.sh
@@ -0,0 +1,15 @@
+# Default PRIMA Stage 1 inference checkpoint:
+#   data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt
+#
+# If this local file is missing, it will be downloaded from the PRIMA Hugging Face repo.
+checkpoint='data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt'
+
+# Update this to your video path before running.
+video_path='demo_data/hati.mp4'
+
+python3 demo_video.py \
+  --checkpoint "${checkpoint}" \
+  --video_path "${video_path}" \
+  --out_video demo_video_out.mp4 \
+  --rotate auto \
+  --full_frame
diff --git a/eval.py b/eval.py
index b7031c5..659b7eb 100644
--- a/eval.py
+++ b/eval.py
@@ -25,6 +25,7 @@ def main(args):
     default_cfg = get_config(args.default_eval_config)
     model = PRIMA.load_from_checkpoint(args.checkpoint, cfg=cfg, strict=False)
     model.eval()
+    model = model.to(args.device)
 
     smal_evaluator = Evaluator(smal_model=model.smal, image_size=cfg.MODEL.IMAGE_SIZE)
     cfg_eval_dataset = dict(default_cfg.DATASETS)
diff --git a/images/teaser.png b/images/teaser.png
index 1f0a3b8..0565e56 100644
Binary files a/images/teaser.png and b/images/teaser.png differ
diff --git a/packages.txt b/packages.txt
new file mode 100644
index 0000000..aaca018
--- /dev/null
+++ b/packages.txt
@@ -0,0 +1,7 @@
+libosmesa6
+libgl1
+libgl1-mesa-dri
+libegl-mesa0
+libegl1
+libglx-mesa0
+libgles2
diff --git a/prima/models/__init__.py b/prima/models/__init__.py
index 6cc213e..d41a09a 100644
--- a/prima/models/__init__.py
+++ b/prima/models/__init__.py
@@ -42,5 +42,13 @@ def load_prima(checkpoint_path):
         model_cfg.MODEL.BACKBONE.pop('PRETRAINED_WEIGHTS')
         model_cfg.freeze()
 
-    model = PRIMA.load_from_checkpoint(checkpoint_path, strict=False, cfg=model_cfg, map_location='cpu')
+    # Offscreen training renderer is not needed for demo/inference startup and
+    # can fail on some local OpenGL backends.
+    model = PRIMA.load_from_checkpoint(
+        checkpoint_path,
+        strict=False,
+        cfg=model_cfg,
+        map_location='cpu',
+        init_renderer=False,
+    )
     return model, model_cfg
diff --git a/prima/models/heads/__init__.py b/prima/models/heads/__init__.py
index 0c08525..61bc900 100644
--- a/prima/models/heads/__init__.py
+++ b/prima/models/heads/__init__.py
@@ -1 +1,10 @@
+"""
+PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation
+
+Official implementation of the paper:
+"PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation"
+by Xiaohang Yu, Ti Wang, and Mackenzie Weygandt Mathis
+Licensed under a modified MIT license
+"""
+
 from .smal_head import build_smal_head
diff --git a/prima/models/losses.py b/prima/models/losses.py
index 971c0c8..2982d17 100644
--- a/prima/models/losses.py
+++ b/prima/models/losses.py
@@ -11,12 +11,44 @@
 import torch.nn as nn
 import numpy as np
 import pickle
-from pytorch3d.transforms import matrix_to_axis_angle
 import torch.nn.functional as F
 from ..utils.geometry import aa_to_rotmat
 from typing import Dict
 
 
+def matrix_to_axis_angle(rot_mats: torch.Tensor) -> torch.Tensor:
+    """Convert rotation matrices (..., 3, 3) to axis-angle vectors (..., 3).
+
+    This local implementation avoids a hard runtime dependency on PyTorch3D.
+    """
+    if rot_mats.shape[-2:] != (3, 3):
+        raise ValueError(f"Expected (..., 3, 3) rotation matrices, got {rot_mats.shape}")
+
+    trace = rot_mats[..., 0, 0] + rot_mats[..., 1, 1] + rot_mats[..., 2, 2]
+    cos_theta = (trace - 1.0) * 0.5
+    cos_theta = torch.clamp(cos_theta, -1.0, 1.0)
+    theta = torch.acos(cos_theta)
+
+    vee = torch.stack(
+        [
+            rot_mats[..., 2, 1] - rot_mats[..., 1, 2],
+            rot_mats[..., 0, 2] - rot_mats[..., 2, 0],
+            rot_mats[..., 1, 0] - rot_mats[..., 0, 1],
+        ],
+        dim=-1,
+    )
+    sin_theta = torch.sin(theta)
+    eps = 1e-6
+    scale = theta / torch.clamp(2.0 * sin_theta, min=eps)
+    aa = vee * scale.unsqueeze(-1)
+
+    # For very small angles, first-order approximation: aa ~= 0.5 * vee
+    small = theta < 1e-4
+    if small.any():
+        aa = torch.where(small.unsqueeze(-1), 0.5 * vee, aa)
+    return aa
+
+
 
 class DepthLoss(nn.Module):
     """
diff --git a/prima/models/prima.py b/prima/models/prima.py
index f03bcd6..0cb861d 100755
--- a/prima/models/prima.py
+++ b/prima/models/prima.py
@@ -22,8 +22,6 @@
 from ..utils.pylogger import get_pylogger
 from .backbones import create_backbone
 from .heads import build_smal_head
-from ..utils import MeshRenderer
-from ..utils import renderer
 from prima.models.smal_wrapper import SMAL
 from .discriminator import Discriminator
 
@@ -130,6 +128,8 @@ def __init__(self, cfg: CfgNode, init_renderer: bool = True):
         # init depth renderer for supervised training 
         # Setup renderer for visualization
         if init_renderer:
+            from ..utils import MeshRenderer
+
             self.mesh_renderer = MeshRenderer(self.cfg, faces=self.smal.faces.numpy())
         else:
             self.mesh_renderer = None
@@ -612,4 +612,4 @@ def validation_step(self, batch: Dict, batch_idx: int, dataloader_idx=0) -> Dict
             # Use global_step for step count when logging validation visuals
             self.tensorboard_logging(batch, output, self.global_step, train=False)
 
-        return output
\ No newline at end of file
+        return output
diff --git a/prima/models/smal_wrapper.py b/prima/models/smal_wrapper.py
index ea9de81..e00d087 100644
--- a/prima/models/smal_wrapper.py
+++ b/prima/models/smal_wrapper.py
@@ -49,13 +49,27 @@ class SMALLayer(nn.Module):
     def __init__(self, num_betas=41, **kwargs):
         super().__init__()
         self.num_betas = num_betas
-        self.register_buffer("shapedirs", torch.from_numpy(np.array(kwargs['shapedirs'], dtype=np.float32))[:, :, :num_betas]) # [3889, 3, 41]
-        self.register_buffer("v_template", torch.from_numpy(np.array(kwargs['v_template']).astype(np.float32)))  # [3889, 3]
-        self.register_buffer("posedirs", torch.from_numpy(np.array(kwargs['posedirs'], dtype=np.float32)).reshape(-1,
-                                                                                                 34*9).T)  # [34*9, 11667]
-        self.register_buffer("J_regressor", torch.from_numpy(kwargs['J_regressor'].toarray().astype(np.float32)))  # [33, 3389]
-        self.register_buffer("lbs_weights", torch.from_numpy(np.array(kwargs['weights'], dtype=np.float32)))  # [3889, 33]
-        self.register_buffer("faces", torch.from_numpy(np.array(kwargs['f'], dtype=np.int32)))  # [7774, 3]
+        from chumpy.ch import materialize
+
+        self.register_buffer(
+            "shapedirs",
+            torch.from_numpy(materialize(kwargs["shapedirs"]))[:, :, :num_betas],
+        )  # [3889, 3, 41]
+        self.register_buffer(
+            "v_template", torch.from_numpy(materialize(kwargs["v_template"]))
+        )  # [3889, 3]
+        self.register_buffer(
+            "posedirs",
+            torch.from_numpy(materialize(kwargs["posedirs"])).reshape(-1, 34 * 9).T,
+        )  # [34*9, 11667]
+        self.register_buffer(
+            "J_regressor",
+            torch.from_numpy(kwargs["J_regressor"].toarray().astype(np.float32)),
+        )  # [33, 3389]
+        self.register_buffer(
+            "lbs_weights", torch.from_numpy(materialize(kwargs["weights"]))
+        )  # [3889, 33]
+        self.register_buffer("faces", torch.from_numpy(materialize(kwargs["f"], dtype=np.int32)))  # [7774, 3]
 
         kintree_table = kwargs['kintree_table']
         self.register_buffer("parents", torch.from_numpy(kintree_table[0].astype(np.int32)))
diff --git a/prima/utils/__init__.py b/prima/utils/__init__.py
index 5a5ab0b..9b323a5 100755
--- a/prima/utils/__init__.py
+++ b/prima/utils/__init__.py
@@ -7,12 +7,10 @@
 Licensed under a modified MIT license
 """
 
-import torch
 from typing import Any
-from .mesh_renderer import MeshRenderer
 
 
-def recursive_to(x: Any, target: torch.device):
+def recursive_to(x: Any, target: Any):
     """
     Recursively transfer a batch of data to the target device
     Args:
@@ -21,11 +19,27 @@ def recursive_to(x: Any, target: torch.device):
     Returns:
         Batch of data where all tensors are transferred to the target device.
     """
-    if isinstance(x, dict):
-        return {k: recursive_to(v, target) for k, v in x.items()}
-    elif isinstance(x, torch.Tensor):
-        return x.to(target)
-    elif isinstance(x, list):
-        return [recursive_to(i, target) for i in x]
-    else:
-        return x
\ No newline at end of file
+    import torch
+
+    def move(value: Any):
+        if isinstance(value, dict):
+            return {k: move(v) for k, v in value.items()}
+        if isinstance(value, torch.Tensor):
+            return value.to(target)
+        if isinstance(value, list):
+            return [move(i) for i in value]
+        return value
+
+    return move(x)
+
+
+def __getattr__(name: str):
+    if name == "MeshRenderer":
+        from .mesh_renderer import MeshRenderer
+
+        globals()[name] = MeshRenderer
+        return MeshRenderer
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = ["MeshRenderer", "recursive_to"]
diff --git a/prima/utils/detection.py b/prima/utils/detection.py
new file mode 100644
index 0000000..72dcae2
--- /dev/null
+++ b/prima/utils/detection.py
@@ -0,0 +1,118 @@
+"""
+PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation
+
+Official implementation of the paper:
+"PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation"
+by Xiaohang Yu, Ti Wang, and Mackenzie Weygandt Mathis
+Licensed under a modified MIT license
+"""
+
+from __future__ import annotations
+
+# Utilities for filtering animal detections before PRIMA demo inference.
+#
+# Detectron2 may return both a full-animal box and a local/partial box for the
+# same animal. These helpers keep the demo pipeline from rendering the same
+# animal multiple times.
+
+from typing import Iterable
+
+import numpy as np
+
+ANIMAL_COCO_IDS = (15, 16, 17, 18, 19, 21, 22)
+
+
+def _box_areas(boxes: np.ndarray) -> np.ndarray:
+    widths = np.maximum(0.0, boxes[:, 2] - boxes[:, 0])
+    heights = np.maximum(0.0, boxes[:, 3] - boxes[:, 1])
+    return widths * heights
+
+
+def _intersection_areas(box: np.ndarray, boxes: np.ndarray) -> np.ndarray:
+    x1 = np.maximum(box[0], boxes[:, 0])
+    y1 = np.maximum(box[1], boxes[:, 1])
+    x2 = np.minimum(box[2], boxes[:, 2])
+    y2 = np.minimum(box[3], boxes[:, 3])
+    return np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1)
+
+
+def _suppress_duplicate_boxes(
+    boxes: np.ndarray,
+    scores: np.ndarray,
+    *,
+    iou_threshold: float,
+    containment_threshold: float,
+) -> np.ndarray:
+    if len(boxes) <= 1:
+        return np.arange(len(boxes), dtype=np.int64)
+
+    boxes = boxes.astype(np.float32, copy=False)
+    scores = scores.astype(np.float32, copy=False)
+    areas = _box_areas(boxes)
+
+    contained = np.zeros(len(boxes), dtype=bool)
+    for idx, area in enumerate(areas):
+        if area <= 0:
+            contained[idx] = True
+            continue
+        larger = np.where(areas > area)[0]
+        if len(larger) == 0:
+            continue
+        covered = _intersection_areas(boxes[idx], boxes[larger]) / area
+        if np.any(covered >= containment_threshold):
+            contained[idx] = True
+
+    candidates = np.where(~contained)[0]
+    if len(candidates) <= 1:
+        return candidates
+
+    order = candidates[np.argsort(scores[candidates])[::-1]]
+    keep = []
+    while len(order) > 0:
+        current = order[0]
+        keep.append(current)
+        rest = order[1:]
+        if len(rest) == 0:
+            break
+
+        inter = _intersection_areas(boxes[current], boxes[rest])
+        union = areas[current] + areas[rest] - inter
+        iou = np.divide(inter, union, out=np.zeros_like(inter), where=union > 0)
+        order = rest[iou <= iou_threshold]
+
+    return np.array(sorted(keep), dtype=np.int64)
+
+
+def select_animal_boxes(
+    det_instances,
+    *,
+    animal_class_ids: Iterable[int] = ANIMAL_COCO_IDS,
+    score_threshold: float = 0.7,
+    iou_threshold: float = 0.5,
+    containment_threshold: float = 0.9,
+) -> tuple[np.ndarray, int]:
+    """Return filtered animal boxes and the number of duplicate boxes removed."""
+    class_ids = set(int(class_id) for class_id in animal_class_ids)
+    classes = det_instances.pred_classes.detach().cpu().numpy()
+    scores = det_instances.scores.detach().cpu().numpy()
+
+    valid_idx = np.array(
+        [
+            i
+            for i, (class_id, score) in enumerate(zip(classes, scores))
+            if int(class_id) in class_ids and float(score) > float(score_threshold)
+        ],
+        dtype=np.int64,
+    )
+    if len(valid_idx) == 0:
+        return np.zeros((0, 4), dtype=np.float32), 0
+
+    boxes = det_instances.pred_boxes.tensor[valid_idx].detach().cpu().numpy()
+    scores = scores[valid_idx]
+    keep = _suppress_duplicate_boxes(
+        boxes,
+        scores,
+        iou_threshold=iou_threshold,
+        containment_threshold=containment_threshold,
+    )
+    return boxes[keep], int(len(boxes) - len(keep))
diff --git a/prima/utils/evaluate_metric.py b/prima/utils/evaluate_metric.py
index d57580f..39720bd 100644
--- a/prima/utils/evaluate_metric.py
+++ b/prima/utils/evaluate_metric.py
@@ -109,34 +109,29 @@ def __init__(self, smal_model, image_size: int=256, pelvis_ind: int = 7):
         self.image_size = image_size
     
     def compute_pck(self, output: Dict, batch: Dict, pck_threshold: Union[List, None]):
+        if pck_threshold is None or len(pck_threshold) == 0:
+            return torch.tensor([], dtype=torch.float32)
+
         pred_keypoints_2d = output['pred_keypoints_2d'].detach().cpu()
         gt_keypoints_2d = batch['keypoints_2d'].detach().cpu()
-        self.pck_threshold_list = []
-        
-        pred_keypoints_2d = (pred_keypoints_2d + 0.5) * self.image_size  # * batch['bbox_expand_factor'].detach().cpu().numpy().reshape(-1, 1, 1)
+
+        pred_keypoints_2d = (pred_keypoints_2d + 0.5) * self.image_size
         conf = gt_keypoints_2d[:, :, -1]
-        gt_keypoints_2d = (gt_keypoints_2d[:, :, :-1] + 0.5) * self.image_size  # * batch['bbox_expand_factor'].detach().cpu().numpy().reshape(-1, 1, 1)
-        if pck_threshold is not None:
-            for i in range(len(pck_threshold)):
-                self.pck_threshold_list.append(torch.tensor([pck_threshold[i]] * len(pred_keypoints_2d), dtype=torch.float32))
-        if len(self.pck_threshold_list) == 0:
-            return torch.tensor([], dtype=torch.float32)
+        gt_keypoints_2d = (gt_keypoints_2d[:, :, :-1] + 0.5) * self.image_size
 
-        pcks = []
-        # Use mask area if available, otherwise use full image area
         if 'mask' in batch and batch['mask'] is not None:
             seg_area = torch.sum(batch['mask'].detach().cpu().reshape(batch['mask'].shape[0], -1), dim=-1).unsqueeze(-1)
         else:
-            # Use full image area as fallback
             seg_area = torch.tensor([self.image_size * self.image_size] * len(pred_keypoints_2d), dtype=torch.float32).unsqueeze(-1)
-        total_visible = torch.sum(conf, dim=-1).clamp_min(1e-6)
-        for th in self.pck_threshold_list:
-            dist = torch.norm(pred_keypoints_2d - gt_keypoints_2d, dim=-1)
 
-            hits = (dist / torch.sqrt(seg_area)) < th.unsqueeze(1)
-            pck = torch.sum(hits.float() * conf, dim=-1) / total_visible
-            pcks.append(pck.numpy().tolist())
-        return torch.mean(torch.tensor(pcks), dim=1)
+        total_visible = torch.sum(conf, dim=-1).clamp_min(1e-6)                # (B,)
+        dist = torch.norm(pred_keypoints_2d - gt_keypoints_2d, dim=-1)         # (B, K)
+        norm_dist = dist / torch.sqrt(seg_area)                                # (B, K)
+
+        thresholds = torch.tensor(pck_threshold, dtype=torch.float32).view(-1, 1, 1)  # (T, 1, 1)
+        hits = (norm_dist.unsqueeze(0) < thresholds).float()                   # (T, B, K)
+        pcks = (hits * conf.unsqueeze(0)).sum(dim=-1) / total_visible.unsqueeze(0)    # (T, B)
+        return pcks.mean(dim=1)                                                # (T,)
 
     def compute_pa_mpjpe(self, pred_joints, gt_joints):
         S1_hat = compute_similarity_transform(pred_joints, gt_joints)
@@ -181,15 +176,11 @@ def eval_2d(self, output: Dict, batch: Dict, pck_threshold: List[float]=[0.10, 0
         auc = self.compute_auc(batch, output)
         return pck.tolist(), auc
     
-    def compute_auc(self, batch: Dict, output: Dict, threshold_min: int=0.0, threshold_max: int=1.0, steps: int=100):
+    def compute_auc(self, batch: Dict, output: Dict, threshold_min: float=0.0, threshold_max: float=1.0, steps: int=100):
         thresholds = np.linspace(threshold_min, threshold_max, steps)
-        norm_factor = np.trapz(np.ones_like(thresholds), thresholds)
-        pck_curve = []
-        for th in thresholds:
-            pck_curve.append(self.compute_pck(output, batch, [th]))
-        pck_curve = torch.tensor(pck_curve).tolist()
-        auc = np.trapz(pck_curve, thresholds)
-        auc /= norm_factor
+        pck_curve = self.compute_pck(output, batch, thresholds.tolist()).numpy()  # (steps,)
+        norm_factor = threshold_max - threshold_min
+        auc = float(np.trapz(pck_curve, thresholds) / norm_factor)
         return auc
 
     def smal_forward(self, batch: Dict):
diff --git a/prima/utils/mesh_renderer.py b/prima/utils/mesh_renderer.py
index df63c79..7501004 100644
--- a/prima/utils/mesh_renderer.py
+++ b/prima/utils/mesh_renderer.py
@@ -8,10 +8,13 @@
 """
 
 import os
+from ctypes.util import find_library
 
-if 'PYOPENGL_PLATFORM' not in os.environ:
-    # EGL is usually unavailable on macOS; use pyglet there.
-    os.environ['PYOPENGL_PLATFORM'] = 'pyglet' if os.uname().sysname == 'Darwin' else 'egl'
+if 'PYOPENGL_PLATFORM' not in os.environ and os.uname().sysname != 'Darwin':
+    # Prefer EGL; PyOpenGL's OSMesa bindings can lack symbols required by pyrender.
+    os.environ['PYOPENGL_PLATFORM'] = 'egl' if find_library('EGL') else 'osmesa'
+    if os.environ['PYOPENGL_PLATFORM'] == 'egl':
+        os.environ.setdefault('EGL_PLATFORM', 'surfaceless')
 import torch
 from torchvision.utils import make_grid
 import numpy as np
diff --git a/prima/utils/renderer.py b/prima/utils/renderer.py
index 9af2824..1c564f8 100644
--- a/prima/utils/renderer.py
+++ b/prima/utils/renderer.py
@@ -10,10 +10,13 @@
 
 
 import os
+from ctypes.util import find_library
 
-if 'PYOPENGL_PLATFORM' not in os.environ:
-    # EGL is usually unavailable on macOS; use pyglet there.
-    os.environ['PYOPENGL_PLATFORM'] = 'pyglet' if os.uname().sysname == 'Darwin' else 'egl'
+if 'PYOPENGL_PLATFORM' not in os.environ and os.uname().sysname != 'Darwin':
+    # Prefer EGL; PyOpenGL's OSMesa bindings can lack symbols required by pyrender.
+    os.environ['PYOPENGL_PLATFORM'] = 'egl' if find_library('EGL') else 'osmesa'
+    if os.environ['PYOPENGL_PLATFORM'] == 'egl':
+        os.environ.setdefault('EGL_PLATFORM', 'surfaceless')
 import torch
 import numpy as np
 import pyrender
@@ -193,9 +196,13 @@ def __call__(self,
         """
 
         if full_frame:
-            
-            image = cv2.imread(imgname)
-            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.
+            if imgname is not None:
+                image = cv2.imread(imgname)
+                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.
+            else:
+                image = image.astype(np.float32)
+                if image.max() > 1.0:
+                    image = image / 255.0
         else:
             image = (image.clone()) * (torch.tensor(self.cfg.MODEL.IMAGE_STD, device=image.device).reshape(3, 1, 1))
             image = image + torch.tensor(self.cfg.MODEL.IMAGE_MEAN, device=image.device).reshape(3, 1, 1)
@@ -204,9 +211,18 @@ def __call__(self,
         # Use custom focal length if provided, otherwise use default
         focal_length_to_use = focal_length if focal_length is not None else self.focal_length
         
-        renderer = pyrender.OffscreenRenderer(viewport_width=image.shape[1],
-                                              viewport_height=image.shape[0],
-                                              point_size=1.0)
+        try:
+            renderer = pyrender.OffscreenRenderer(
+                viewport_width=image.shape[1],
+                viewport_height=image.shape[0],
+                point_size=1.0,
+            )
+        except (IndexError, OSError) as exc:
+            raise RuntimeError(
+                "PyRender could not open an OpenGL context (common on headless macOS or remote SSH). "
+                "Run the demo from a normal desktop session, or on Linux/Spaces use OSMesa (see packages.txt). "
+                f"Original error: {exc}"
+            ) from exc
         material = pyrender.MetallicRoughnessMaterial(
             metallicFactor=0.0,
             alphaMode='OPAQUE',
@@ -250,7 +266,7 @@ def __call__(self,
         if return_rgba:
             return color
 
-        valid_mask = (color[:, :, -1])[:, :, np.newaxis]
+        valid_mask = (rend_depth > 0).astype(np.float32)[:, :, np.newaxis]
         if not side_view:
             output_img = (color[:, :, :3] * valid_mask + (1 - valid_mask) * image)
         else:
@@ -428,6 +444,3 @@ def add_point_lighting(self, scene, cam_node, color=np.ones(3), intensity=1.0):
             if scene.has_node(node):
                 continue
             scene.add_node(node)
-
-
-
diff --git a/prima/utils/weights.py b/prima/utils/weights.py
new file mode 100644
index 0000000..c02e593
--- /dev/null
+++ b/prima/utils/weights.py
@@ -0,0 +1,337 @@
+"""
+PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation
+
+Official implementation of the paper:
+"PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation"
+by Xiaohang Yu, Ti Wang, and Mackenzie Weygandt Mathis
+Licensed under a modified MIT license
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+from pathlib import Path
+from typing import Iterable, Optional, Sequence, Union
+
+HF_REPO_ID = "MLAdaptiveIntelligence/PRIMA"
+DEFAULT_HF_REPO_ID = HF_REPO_ID
+
+DEFAULT_STAGE1_CHECKPOINT = Path("data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt")
+DEFAULT_STAGE3_CHECKPOINT = Path("data/PRIMAS3/checkpoints/s3ckpt_inference.ckpt")
+
+SMAL_ASSET_PATHS = [
+    "my_smpl_00781_4_all.pkl",
+    "my_smpl_data_00781_4_all.pkl",
+    "walking_toy_symmetric_pose_prior_with_cov_35parts.pkl",
+]
+BACKBONE_ASSET_PATH = "amr_vitbb.pth"
+STAGE1_CONFIG_ASSET_PATH = "config_s1_HYDRA.yaml"
+STAGE1_CHECKPOINT_ASSET_PATH = "s1ckpt_inference.ckpt"
+STAGE3_CONFIG_ASSET_PATH = "config_s3_HYDRA.yaml"
+STAGE3_CHECKPOINT_ASSET_PATH = "s3ckpt_inference.ckpt"
+
+STAGE_ASSETS = {
+    "PRIMAS1": (STAGE1_CONFIG_ASSET_PATH, STAGE1_CHECKPOINT_ASSET_PATH, "s1ckpt_inference.ckpt"),
+    "PRIMAS3": (STAGE3_CONFIG_ASSET_PATH, STAGE3_CHECKPOINT_ASSET_PATH, "s3ckpt_inference.ckpt"),
+}
+
+STAGE_CHECKPOINTS = {
+    "PRIMAS1": Path("PRIMAS1/checkpoints/s1ckpt_inference.ckpt"),
+    "PRIMAS3": Path("PRIMAS3/checkpoints/s3ckpt_inference.ckpt"),
+}
+
+PathLike = Union[str, Path]
+
+
+def _resolve_hf_repo_id(hf_repo_id: Optional[str]) -> str:
+    return hf_repo_id or os.environ.get("PRIMA_HF_REPO_ID", HF_REPO_ID)
+
+
+def _default_checkpoint_path(data_dir: PathLike = "data") -> Path:
+    return Path(data_dir) / STAGE_CHECKPOINTS["PRIMAS1"]
+
+
+def _config_path_for_checkpoint(checkpoint_path: PathLike) -> Path:
+    checkpoint_path = Path(checkpoint_path)
+    return checkpoint_path.parent.parent / ".hydra" / "config.yaml"
+
+
+def _stage_for_checkpoint(checkpoint_path: PathLike) -> Optional[str]:
+    checkpoint_path = Path(checkpoint_path)
+    if len(checkpoint_path.parents) < 2:
+        return None
+    stage_name = checkpoint_path.parent.parent.name
+    stage_assets = STAGE_ASSETS.get(stage_name)
+    if stage_assets is None:
+        return None
+    _, _, checkpoint_name = stage_assets
+    if checkpoint_path.name != checkpoint_name:
+        return None
+    return stage_name
+
+
+def _download_file(
+    hf_repo_id: str,
+    remote_filename: str,
+    destination: Path,
+    force_download: bool = False,
+) -> None:
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError:
+        raise ImportError(
+            "huggingface_hub is required to download PRIMA demo assets. "
+            "Install it with:  pip install huggingface_hub\n"
+            "Or download the assets manually and pass a local checkpoint path."
+        ) from None
+
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    downloaded = hf_hub_download(
+        repo_id=hf_repo_id,
+        filename=remote_filename,
+        local_dir=str(destination.parent),
+        local_dir_use_symlinks=False,
+        force_download=force_download,
+    )
+    downloaded_path = Path(downloaded).resolve()
+    target = destination.resolve()
+    if downloaded_path != target:
+        if target.exists():
+            target.unlink()
+        shutil.move(str(downloaded_path), str(target))
+
+
+def _validate_torch_checkpoint(path: Path) -> None:
+    import inspect
+    import pickle
+    import zipfile
+
+    import torch
+
+    if zipfile.is_zipfile(path):
+        with zipfile.ZipFile(path) as checkpoint_zip:
+            corrupt_member = checkpoint_zip.testzip()
+        if corrupt_member is not None:
+            raise RuntimeError(
+                f"Checkpoint file is invalid or incomplete: {path}\n"
+                f"Corrupt archive member: {corrupt_member}\n"
+                "Please redownload the checkpoint and try again."
+            )
+
+    supports_weights_only = "weights_only" in inspect.signature(torch.load).parameters
+    load_kwargs = {"map_location": "cpu"}
+    if supports_weights_only:
+        load_kwargs["weights_only"] = True
+
+    try:
+        torch.load(path, **load_kwargs)
+    except pickle.UnpicklingError as exc:
+        message = str(exc)
+        if (
+            supports_weights_only
+            and "Weights only load failed" in message
+            and ("Unsupported global" in message or "Unsupported class" in message)
+        ):
+            return
+        raise RuntimeError(
+            f"Checkpoint file is invalid or incomplete: {path}\n"
+            "Downloaded checkpoint is not loadable. "
+            "Please verify the uploaded Hugging Face file and try again."
+        ) from exc
+    except Exception as exc:
+        raise RuntimeError(
+            f"Checkpoint file is invalid or incomplete: {path}\n"
+            "Downloaded checkpoint is not loadable. "
+            "Please verify the uploaded Hugging Face file and try again."
+        ) from exc
+
+
+def _ensure_backbone(data_dir: Path, force: bool, hf_repo_id: str) -> None:
+    target = data_dir / "amr_vitbb.pth"
+    if target.exists() and not force:
+        print(f"[skip] {target} already exists")
+        return
+
+    print("[download] pretrained backbone")
+    _download_file(hf_repo_id, BACKBONE_ASSET_PATH, target, force_download=force)
+    print(f"[ok] {target}")
+
+
+def _ensure_smal_assets(data_dir: Path, force: bool, hf_repo_id: str) -> None:
+    required = [Path(p).name for p in SMAL_ASSET_PATHS]
+    smal_dir = data_dir / "smal"
+    if smal_dir.exists() and all((smal_dir / n).exists() for n in required) and not force:
+        print("[skip] SMAL files already exist")
+        return
+
+    print("[download] SMAL assets")
+    for asset_path in SMAL_ASSET_PATHS:
+        target = smal_dir / Path(asset_path).name
+        _download_file(hf_repo_id, asset_path, target, force_download=force)
+    print(f"[ok] {smal_dir}")
+
+
+def _ensure_stage_assets(
+    stage_name: str,
+    data_dir: Path,
+    force: bool,
+    hf_repo_id: str,
+    validate_existing: bool = True,
+) -> None:
+    if stage_name not in STAGE_ASSETS:
+        known = ", ".join(sorted(STAGE_ASSETS))
+        raise ValueError(f"Unknown PRIMA stage '{stage_name}'. Expected one of: {known}")
+
+    config_asset_path, checkpoint_asset_path, checkpoint_name = STAGE_ASSETS[stage_name]
+    stage_dir = data_dir / stage_name
+    config_target = stage_dir / ".hydra" / "config.yaml"
+    checkpoint_target = stage_dir / "checkpoints" / checkpoint_name
+    redownload_checkpoint = False
+
+    if config_target.exists() and checkpoint_target.exists() and not force:
+        if validate_existing:
+            try:
+                _validate_torch_checkpoint(checkpoint_target)
+            except RuntimeError:
+                print(f"[warn] {stage_name} checkpoint is incomplete, redownloading checkpoint only.")
+                redownload_checkpoint = True
+            else:
+                print(f"[skip] {stage_name} assets already exist")
+                return
+        else:
+            print(f"[skip] {stage_name} assets already exist")
+            return
+
+    print(f"[download] {stage_name} assets")
+    config_target.parent.mkdir(parents=True, exist_ok=True)
+    checkpoint_target.parent.mkdir(parents=True, exist_ok=True)
+    if force or not config_target.exists():
+        _download_file(hf_repo_id, config_asset_path, config_target, force_download=force)
+    if redownload_checkpoint and checkpoint_target.exists():
+        checkpoint_target.unlink()
+    if force or redownload_checkpoint or not checkpoint_target.exists():
+        _download_file(
+            hf_repo_id,
+            checkpoint_asset_path,
+            checkpoint_target,
+            force_download=force or redownload_checkpoint,
+        )
+    _validate_torch_checkpoint(checkpoint_target)
+    print(f"[ok] {stage_dir}")
+
+
+def _normalize_stages(stages: Union[str, Iterable[str]]) -> Sequence[str]:
+    if isinstance(stages, str):
+        return (stages,)
+    return tuple(stages)
+
+
+def _verify_assets(data_dir: Path, stages: Sequence[str]) -> None:
+    required_paths = [
+        data_dir / "smal" / "my_smpl_00781_4_all.pkl",
+        data_dir / "smal" / "my_smpl_data_00781_4_all.pkl",
+        data_dir / "smal" / "walking_toy_symmetric_pose_prior_with_cov_35parts.pkl",
+        data_dir / "amr_vitbb.pth",
+    ]
+    for stage_name in stages:
+        if stage_name not in STAGE_ASSETS:
+            known = ", ".join(sorted(STAGE_ASSETS))
+            raise ValueError(f"Unknown PRIMA stage '{stage_name}'. Expected one of: {known}")
+        _, _, checkpoint_name = STAGE_ASSETS[stage_name]
+        stage_dir = data_dir / stage_name
+        required_paths.extend(
+            [
+                stage_dir / ".hydra" / "config.yaml",
+                stage_dir / "checkpoints" / checkpoint_name,
+            ]
+        )
+
+    missing = [p for p in required_paths if not p.exists()]
+    if missing:
+        raise FileNotFoundError("Missing required files:\n" + "\n".join(str(p) for p in missing))
+
+    for stage_name in stages:
+        _, _, checkpoint_name = STAGE_ASSETS[stage_name]
+        _validate_torch_checkpoint(data_dir / stage_name / "checkpoints" / checkpoint_name)
+
+
+def _ensure_assets_for_checkpoint(
+    checkpoint_path: PathLike,
+    force: bool = False,
+    hf_repo_id: Optional[str] = None,
+) -> None:
+    checkpoint_path = Path(checkpoint_path)
+    config_path = _config_path_for_checkpoint(checkpoint_path)
+    stage_name = _stage_for_checkpoint(checkpoint_path)
+    if stage_name is None:
+        if checkpoint_path.exists() and config_path.exists() and not force:
+            print(f"[skip] Using local PRIMA checkpoint {checkpoint_path}")
+            return
+        raise FileNotFoundError(
+            "Missing checkpoint or config for a custom path:\n"
+            f"  checkpoint: {checkpoint_path}\n"
+            f"  config: {config_path}\n"
+            "Auto-download supports the standard PRIMA demo layouts only:\n"
+            "  data/PRIMAS1/checkpoints/s1ckpt_inference.ckpt\n"
+            "  data/PRIMAS3/checkpoints/s3ckpt_inference.ckpt\n"
+            "Pass one of those paths, or download/copy your custom checkpoint manually."
+        )
+
+    data_dir = checkpoint_path.parent.parent.parent
+    repo_id = _resolve_hf_repo_id(hf_repo_id)
+    print(f"[download] Ensuring PRIMA demo assets under {data_dir}")
+    _ensure_smal_assets(data_dir, force=force, hf_repo_id=repo_id)
+    _ensure_backbone(data_dir, force=force, hf_repo_id=repo_id)
+    _ensure_stage_assets(
+        stage_name,
+        data_dir,
+        force=force,
+        hf_repo_id=repo_id,
+        validate_existing=False,
+    )
+
+
+def ensure_demo_assets(
+    data_dir: PathLike = "data",
+    *,
+    stages: Union[str, Iterable[str]] = ("PRIMAS1",),
+    force: bool = False,
+    hf_repo_id: Optional[str] = None,
+) -> None:
+    """Ensure PRIMA demo assets exist in the expected ``data/`` layout."""
+    data_dir = Path(data_dir).resolve()
+    data_dir.mkdir(parents=True, exist_ok=True)
+    repo_id = _resolve_hf_repo_id(hf_repo_id)
+    selected_stages = _normalize_stages(stages)
+
+    _ensure_smal_assets(data_dir, force=force, hf_repo_id=repo_id)
+    _ensure_backbone(data_dir, force=force, hf_repo_id=repo_id)
+    for stage_name in selected_stages:
+        _ensure_stage_assets(stage_name, data_dir, force=force, hf_repo_id=repo_id)
+    _verify_assets(data_dir, selected_stages)
+
+
+def resolve_prima_checkpoint_path(
+    checkpoint_path: PathLike = "",
+    *,
+    data_dir: PathLike = "data",
+    auto_download: bool = True,
+    hf_repo_id: Optional[str] = None,
+    force: bool = False,
+) -> str:
+    """Return a PRIMA checkpoint path, downloading default demo assets if needed."""
+    resolved = Path(checkpoint_path) if checkpoint_path else _default_checkpoint_path(data_dir)
+    if auto_download:
+        _ensure_assets_for_checkpoint(resolved, force=force, hf_repo_id=hf_repo_id)
+    return str(resolved)
+
+
+__all__ = [
+    "DEFAULT_HF_REPO_ID",
+    "DEFAULT_STAGE1_CHECKPOINT",
+    "DEFAULT_STAGE3_CHECKPOINT",
+    "HF_REPO_ID",
+    "ensure_demo_assets",
+    "resolve_prima_checkpoint_path",
+]
diff --git a/pyproject.toml b/pyproject.toml
index 4b1de79..17576dd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,7 +7,7 @@ name = "prima-animal"
 version = "0.1.7"
 description = "PRIMA: 3D animal pose and shape estimation"
 readme = "README.md"
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 
 authors = [
   { name = "Xiaohang Yu", email = "xiaohang.yu@epfl.ch" },
@@ -57,13 +57,9 @@ dependencies = [
 
   # IO / misc
   "gdown==5.2.0",
-  "gradio==5.1.0",
-  "pydantic==2.10.6",
-
-  # Web stack for Gradio demo
-  "fastapi==0.110.0",
-  "starlette==0.36.3",
-  "jinja2==3.1.4",
+  # Match HF Space (Gradio 6.x) and local demo; do not pin 5.1 — Space injects gradio==6.x.
+  "gradio>=5.1,<7",
+  "pydantic>=2.10,<3",
 
   # Training framework
   "pytorch-lightning==2.5.5",
@@ -76,7 +72,7 @@ dependencies = [
 
   # Demo runtime dependencies (included in main PyPI install)
   "detectron2 @ git+https://github.com/facebookresearch/detectron2.git",
-  "deeplabcut",
+  "deeplabcut==3.0.0rc14",
 ]
 
 [project.optional-dependencies]
@@ -92,4 +88,4 @@ include-package-data = true
 
 [tool.setuptools.packages.find]
 where = ["."]
-include = ["prima*"]
+include = ["prima*", "chumpy*"]
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..b6fc543
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,38 @@
+huggingface_hub<1
+torch==2.2.1
+torchvision==0.17.1
+# Do not pin gradio here: Hugging Face Spaces runs
+#   pip install -r requirements.txt gradio[oauth,mcp]==6.x ...
+# and a pinned older gradio causes ResolutionImpossible. Local installs get Gradio from
+# scripts/clean_install_local.sh after this file, or from pyproject.toml for editable installs.
+# 3.x rc needed for ``pose_estimation_pytorch`` / SuperAnimal (2.3.x is TensorFlow-only).
+deeplabcut==3.0.0rc14
+# PyTables: use a wheel on macOS (DLC 2.x pinned 3.8.0 from source); Linux pip resolves normally.
+tables>=3.9.2,<3.11
+# Animal detector (needs torch installed first; local clean_install uses --no-build-isolation).
+detectron2 @ git+https://github.com/facebookresearch/detectron2.git
+pytorch-lightning==2.5.5
+yacs==0.1.8
+pyrender==0.1.45
+trimesh==4.8.2
+opencv-python==4.11.0.86
+timm==1.0.24
+einops==0.8.1
+smplx==0.1.28
+xtcocotools==1.14.3
+open_clip_torch==3.2.0
+transformers==4.56.2
+omegaconf==2.3.0
+hydra-core==1.3.2
+hydra-submitit-launcher==1.2.0
+hydra-colorlog==1.2.0
+pyrootutils==1.0.4
+rich==14.1.0
+scikit-image==0.25.2
+pandas==2.3.2
+numpy==1.26.1
+gdown==5.2.0
+setuptools<81
+packaging<25
+Cython<3
+wheel
diff --git a/scripts/clean_install_local.sh b/scripts/clean_install_local.sh
new file mode 100755
index 0000000..b6fab94
--- /dev/null
+++ b/scripts/clean_install_local.sh
@@ -0,0 +1,230 @@
+#!/usr/bin/env bash
+# Fresh local environment: venv, pip deps, LFS assets, demo checkpoints, smoke test.
+#
+# Requires Python 3.10+ (matches README, Space, and type hints in app.py).
+#
+# Usage:
+#   ./scripts/clean_install_local.sh
+#   PRIMA_PYTHON=/opt/homebrew/bin/python3.10 ./scripts/clean_install_local.sh
+#   PRIMA_VENV=.venv ./scripts/clean_install_local.sh --skip-data
+#   ./scripts/clean_install_local.sh --wipe-data --force-data
+set -euo pipefail
+
+# Non-interactive: no pip/git credential prompts on stdin.
+export GIT_TERMINAL_PROMPT=0
+export PIP_DISABLE_PIP_VERSION_CHECK=1
+export HF_HUB_DISABLE_SYMLINKS_WARNING=1
+
+ROOT="$(git rev-parse --show-toplevel)"
+cd "$ROOT"
+
+VENV="${PRIMA_VENV:-.venv}"
+SKIP_DATA=0
+FORCE_DATA=0
+WIPE_DATA=0
+EDITABLE=1
+
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --venv)
+      VENV="$2"
+      shift 2
+      ;;
+    --skip-data)
+      SKIP_DATA=1
+      shift
+      ;;
+    --force-data)
+      FORCE_DATA=1
+      shift
+      ;;
+    --wipe-data)
+      WIPE_DATA=1
+      shift
+      ;;
+    --no-editable)
+      EDITABLE=0
+      shift
+      ;;
+    -h|--help)
+      echo "Usage: $0 [--venv DIR] [--skip-data] [--force-data] [--wipe-data] [--no-editable]"
+      echo "Env: PRIMA_PYTHON=python3.10  PRIMA_VENV=.venv"
+      exit 0
+      ;;
+    *)
+      echo "Unknown option: $1" >&2
+      exit 1
+      ;;
+  esac
+done
+
+resolve_python() {
+  if [[ -n "${PRIMA_PYTHON:-}" ]]; then
+    if [[ -x "${PRIMA_PYTHON}" ]] || command -v "${PRIMA_PYTHON}" >/dev/null 2>&1; then
+      echo "${PRIMA_PYTHON}"
+      return 0
+    fi
+    echo "[clean-install] ERROR: PRIMA_PYTHON=${PRIMA_PYTHON} is not executable." >&2
+    return 1
+  fi
+  local c p
+  for c in python3.10 python3.11; do
+    if command -v "$c" >/dev/null 2>&1; then
+      if "$c" -c 'import sys; raise SystemExit(0 if sys.version_info >= (3, 10) else 1)'; then
+        command -v "$c"
+        return 0
+      fi
+    fi
+  done
+  for p in /opt/homebrew/bin/python3.10 /usr/local/bin/python3.10 /opt/homebrew/opt/python@3.10/bin/python3.10 /usr/local/opt/python@3.10/bin/python3.10; do
+    if [[ -x "$p" ]]; then
+      echo "$p"
+      return 0
+    fi
+  done
+  if command -v python3.12 >/dev/null 2>&1; then
+    command -v python3.12
+    return 0
+  fi
+  return 1
+}
+
+resolve_torch_index_url() {
+  if [[ -n "${PRIMA_TORCH_INDEX_URL:-}" ]]; then
+    echo "${PRIMA_TORCH_INDEX_URL}"
+    return 0
+  fi
+
+  if [[ "$(uname -s)" == "Darwin" ]]; then
+    return 1
+  fi
+
+  if command -v nvcc >/dev/null 2>&1; then
+    local cuda_version
+    cuda_version="$(nvcc --version | sed -n 's/.*release \([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' | head -n 1)"
+    case "${cuda_version}" in
+      11.8)
+        echo "https://download.pytorch.org/whl/cu118"
+        return 0
+        ;;
+      12.1)
+        echo "https://download.pytorch.org/whl/cu121"
+        return 0
+        ;;
+      "")
+        echo "[clean-install] WARN: Could not parse nvcc CUDA version; using pip default PyTorch wheel." >&2
+        return 1
+        ;;
+      *)
+        echo "[clean-install] WARN: CUDA ${cuda_version} detected; set PRIMA_TORCH_INDEX_URL if Detectron2 needs a specific PyTorch wheel." >&2
+        return 1
+        ;;
+    esac
+  fi
+
+  return 1
+}
+
+echo "[clean-install] Repository: ${ROOT}"
+
+if ! PY="$(resolve_python)"; then
+  echo "[clean-install] ERROR: Need Python 3.10, 3.11, or 3.12 on PATH, or set PRIMA_PYTHON." >&2
+  echo "  Conda example: PRIMA_PYTHON=\$HOME/miniconda3/envs/prima/bin/python3.10 $0 ..." >&2
+  echo "  macOS example: brew install python@3.10 && PRIMA_PYTHON=/opt/homebrew/bin/python3.10 $0 ..." >&2
+  exit 1
+fi
+echo "[clean-install] Using Python: $("$PY" -c 'import sys; print(sys.executable, sys.version.split()[0])')"
+
+if command -v git-lfs >/dev/null 2>&1; then
+  echo "[clean-install] git lfs pull (demo images / teaser) ..."
+  git lfs install
+  git lfs pull
+else
+  echo "[clean-install] WARN: git-lfs not found; demo images may be LFS pointer stubs. Install: brew install git-lfs && git lfs install" >&2
+fi
+
+if [[ -d "$VENV" ]]; then
+  echo "[clean-install] Removing existing venv: ${VENV}"
+  rm -rf "$VENV"
+fi
+
+echo "[clean-install] Creating venv: ${VENV}"
+"$PY" -m venv "$VENV"
+# shellcheck disable=SC1090
+source "${VENV}/bin/activate"
+
+python -m pip install --no-input -U pip wheel
+# Match requirements.txt / pyproject pins before pulling the rest
+python -m pip install --no-input "setuptools<81" "packaging<25" "Cython<3"
+python -m pip install --no-input "numpy==1.26.1"
+
+echo "[clean-install] xtcocotools (needs numpy available during build) ..."
+python -m pip install --no-input --no-build-isolation "xtcocotools==1.14.3"
+
+if TORCH_INDEX_URL="$(resolve_torch_index_url)"; then
+  echo "[clean-install] Installing PyTorch from ${TORCH_INDEX_URL} ..."
+  python -m pip install --no-input --index-url "${TORCH_INDEX_URL}" \
+    "torch==2.2.1" "torchvision==0.17.1"
+fi
+
+echo "[clean-install] pip install -r requirements.txt (this can take a long time) ..."
+REQ_TMP="$(mktemp)"
+grep -vE '^[[:space:]]*(deeplabcut|detectron2|xtcocotools)' "${ROOT}/requirements.txt" > "${REQ_TMP}"
+python -m pip install --no-input -r "${REQ_TMP}"
+rm -f "${REQ_TMP}"
+
+if [[ "$(uname -s)" == "Darwin" ]]; then
+  echo "[clean-install] macOS: PyTables wheel then DeepLabCut 3.x (SuperAnimal pytorch API) ..."
+  python -m pip install --no-input "tables>=3.9.2,<3.11"
+  python -m pip install --no-input "deeplabcut==3.0.0rc14" || {
+    echo "[clean-install] ERROR: deeplabcut install failed. Try: brew install hdf5 && retry." >&2
+    exit 1
+  }
+else
+  python -m pip install --no-input "deeplabcut==3.0.0rc14"
+fi
+
+echo "[clean-install] Detectron2 (needs torch in venv; --no-build-isolation) ..."
+python -m pip install --no-input --no-build-isolation \
+  "detectron2 @ git+https://github.com/facebookresearch/detectron2.git"
+
+# Spaces install Gradio separately; local venv needs it for app.py.
+echo "[clean-install] Installing Gradio for local demo (HF Space provides its own) ..."
+python -m pip install --no-input "gradio>=5.1,<7"
+
+if [[ "$EDITABLE" -eq 1 ]]; then
+  echo "[clean-install] pip install --no-deps -e . (register package; runtime deps from requirements.txt) ..."
+  python -m pip install --no-input --no-deps -e "${ROOT}"
+fi
+
+if [[ "$WIPE_DATA" -eq 1 ]]; then
+  echo "[clean-install] Wiping downloaded demo data under data/ ..."
+  rm -rf "${ROOT}/data/PRIMAS1" "${ROOT}/data/PRIMAS3" "${ROOT}/data/smal" "${ROOT}/data/amr_vitbb.pth" 2>/dev/null || true
+fi
+
+if [[ "$SKIP_DATA" -eq 0 ]]; then
+  echo "[clean-install] Downloading demo assets (large) ..."
+  if [[ "$FORCE_DATA" -eq 1 ]]; then
+    python "${ROOT}/scripts/setup_demo_data.py" --force
+  else
+    python "${ROOT}/scripts/setup_demo_data.py"
+  fi
+else
+  echo "[clean-install] Skipping setup_demo_data (--skip-data)."
+fi
+
+export PYTHONPATH="${ROOT}${PYTHONPATH:+:${PYTHONPATH}}"
+
+echo "[clean-install] Smoke test: import app + build_demo + DeepLabCut API ..."
+python -c "
+import app
+app.get_demo_profile.cache_clear()
+p = app.get_demo_profile()
+print('[clean-install] demo profile:', p.mode)
+app.build_demo()
+print('[clean-install] DeepLabCut SuperAnimal (may take ~30s on first import) ...')
+from deeplabcut.pose_estimation_pytorch.apis import superanimal_analyze_images  # noqa: F401
+print('[clean-install] Gradio demo build + DeepLabCut 3.x: OK')
+"
+
+echo "[clean-install] Done. Activate with: source ${VENV}/bin/activate"
diff --git a/scripts/clean_redeploy_hf_space.sh b/scripts/clean_redeploy_hf_space.sh
new file mode 100755
index 0000000..1192a4f
--- /dev/null
+++ b/scripts/clean_redeploy_hf_space.sh
@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+# Clean redeploy of the Hugging Face Space from the current working tree.
+# Same as scripts/deploy_hf_space.sh; use after a local clean install or any code change.
+set -euo pipefail
+ROOT="$(git rev-parse --show-toplevel)"
+exec "${ROOT}/scripts/deploy_hf_space.sh"
diff --git a/scripts/deploy_hf_space.sh b/scripts/deploy_hf_space.sh
new file mode 100755
index 0000000..2c10faa
--- /dev/null
+++ b/scripts/deploy_hf_space.sh
@@ -0,0 +1,145 @@
+#!/usr/bin/env bash
+# Deploy working tree to Hugging Face Space MLAdaptiveIntelligence/PRIMA-demo.
+#
+# Demo PNG/JPG are tracked with Git LFS (Hugging Face Hub Xet bridge); see .gitattributes.
+# We rsync only the Git-tracked files needed by app.py from the working tree
+# (not ``git archive``), so tracked LFS files are materialized bytes while
+# untracked local files and non-Space project files stay out. Then ``git add``
+# stores matching files as LFS objects on push.
+#
+# Prerequisites: brew install git-lfs git-xet && git xet install && git lfs install
+set -euo pipefail
+
+export GIT_TERMINAL_PROMPT=0
+
+ROOT="$(git rev-parse --show-toplevel)"
+cd "$ROOT"
+SPACE_URL="${HF_SPACE_GIT_URL:-https://huggingface.co/spaces/MLAdaptiveIntelligence/PRIMA-demo.git}"
+
+if ! command -v git-lfs >/dev/null 2>&1; then
+  echo "[deploy] ERROR: git-lfs is required. Install: brew install git-lfs && git lfs install" >&2
+  exit 1
+fi
+
+TMP="$(mktemp -d)"
+cleanup() { rm -rf "$TMP"; }
+trap cleanup EXIT
+
+SPACE_SYNC_PATHS=(
+  ".gitattributes"
+  "README.md"
+  "requirements.txt"
+  "pyproject.toml"
+  "app.py"
+  "demo_tta.py"
+  "chumpy"
+  "configs/sa_finetune_hrnet_w32.yaml"
+  "demo_data"
+  "images/teaser.png"
+  "prima"
+)
+SPACE_EXTRA_FILES=(
+  "packages.txt"
+)
+
+echo "[deploy] Rsync Git-tracked Space files from ${ROOT} ..."
+printf '[deploy]   %s\n' "${SPACE_SYNC_PATHS[@]}"
+missing_tracked=()
+for path in "${SPACE_SYNC_PATHS[@]}"; do
+  if [[ -z "$(git ls-files -- "$path")" ]]; then
+    missing_tracked+=("$path")
+  fi
+done
+if [[ "${#missing_tracked[@]}" -gt 0 ]]; then
+  printf '[deploy] ERROR: Space sync path is not tracked by Git: %s\n' "${missing_tracked[@]}" >&2
+  echo "[deploy] Add required new files with git add, or remove them from SPACE_SYNC_PATHS." >&2
+  exit 1
+fi
+git ls-files -z -- "${SPACE_SYNC_PATHS[@]}" | rsync -a --from0 --files-from=- "${ROOT}/" "${TMP}/"
+
+echo "[deploy] Rsync explicit Space config files from ${ROOT} ..."
+for path in "${SPACE_EXTRA_FILES[@]}"; do
+  if [[ ! -f "$path" ]]; then
+    echo "[deploy] ERROR: Missing required Space config file: $path" >&2
+    exit 1
+  fi
+  printf '[deploy]   %s\n' "$path"
+  rsync -a --relative "$path" "$TMP/"
+done
+
+README_FILE="${TMP}/README.md"
+REQ_FILE="${TMP}/requirements.txt"
+
+echo "[deploy] Removing Detectron2 from Space requirements (app falls back to SuperAnimal detection) ..."
+grep -vE '^[[:space:]]*detectron2([[:space:]]|@|$)' "$REQ_FILE" > "${REQ_FILE}.tmp"
+mv "${REQ_FILE}.tmp" "$REQ_FILE"
+
+if ! sed -n '1,20p' "$README_FILE" | grep -q '^sdk: gradio$'; then
+  echo "[deploy] Adding Hugging Face Space YAML front matter to README.md ..."
+  README_TMP="${README_FILE}.tmp"
+  {
+    cat <<'YAML'
+---
+title: PRIMA Demo
+emoji: 🦮
+colorFrom: blue
+colorTo: green
+sdk: gradio
+python_version: "3.10"
+app_file: app.py
+startup_duration_timeout: 60m
+---
+
+YAML
+    cat "$README_FILE"
+  } > "$README_TMP"
+  mv "$README_TMP" "$README_FILE"
+fi
+
+cd "$TMP"
+
+echo "[deploy] Git init + LFS commit ..."
+git init -q
+git lfs install
+git add -A
+git -c user.email="space-deploy@users.noreply.github.com" -c user.name="HF Space deploy" commit -q -m "Deploy snapshot (LFS for demo images per .gitattributes)"
+
+PUSH_URL="$SPACE_URL"
+if [[ "$PUSH_URL" == https://huggingface.co/* && -z "${HF_TOKEN:-}" && -f "${HF_HOME:-$HOME/.cache/huggingface}/token" ]]; then
+  HF_TOKEN="$(<"${HF_HOME:-$HOME/.cache/huggingface}/token")"
+fi
+if [[ "$PUSH_URL" == https://huggingface.co/* && -n "${HF_TOKEN:-}" ]]; then
+  export HF_TOKEN
+  git config credential.helper "store --file=${TMP}/git-credentials"
+  printf 'protocol=https\nhost=huggingface.co\nusername=hf_user\npassword=%s\n\n' "$HF_TOKEN" | git credential approve
+  ASKPASS="${TMP}/git-askpass.sh"
+  cat > "$ASKPASS" <<'SH'
+#!/usr/bin/env bash
+case "$1" in
+  *Username*) printf '%s\n' 'hf_user' ;;
+  *Password*) printf '%s\n' "${HF_TOKEN}" ;;
+  *) printf '%s\n' "${HF_TOKEN}" ;;
+esac
+SH
+  chmod 700 "$ASKPASS"
+  export GIT_ASKPASS="$ASKPASS"
+fi
+
+git remote add hf "$PUSH_URL"
+echo "[deploy] Uploading LFS objects to Hugging Face Space ..."
+mapfile -t LFS_OIDS < <(git lfs ls-files -l | awk '{print $1}')
+if [[ "${#LFS_OIDS[@]}" -gt 0 ]]; then
+  if ! GIT_TERMINAL_PROMPT=0 git lfs push --object-id hf "${LFS_OIDS[@]}"; then
+    echo "[deploy] ERROR: LFS upload failed. Ensure HF_TOKEN has write access to ${SPACE_URL}." >&2
+    exit 1
+  fi
+else
+  echo "[deploy] No LFS objects found in this snapshot."
+fi
+
+echo "[deploy] Force-pushing to Hugging Face Space ..."
+# This deploy repo is freshly initialized, so older git-lfs pre-push hooks can
+# fail when they try to inspect the remote's previous main commit. LFS objects
+# are uploaded explicitly above; skip the hook for the Git ref update.
+GIT_TERMINAL_PROMPT=0 git push hf HEAD:main --force --no-verify
+echo "[deploy] Done."
diff --git a/scripts/local_infer.py b/scripts/local_infer.py
new file mode 100755
index 0000000..1ee1f36
--- /dev/null
+++ b/scripts/local_infer.py
@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+"""
+PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation
+
+Official implementation of the paper:
+"PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation"
+by Xiaohang Yu, Ti Wang, and Mackenzie Weygandt Mathis
+Licensed under a modified MIT license
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+import cv2
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Local PRIMA inference (no Gradio).")
+    p.add_argument(
+        "--image",
+        type=str,
+        default=str(ROOT / "demo_data" / "beagle.jpg"),
+        help="Path to an input image.",
+    )
+    p.add_argument(
+        "--out",
+        type=str,
+        default=str(ROOT / "demo_out_local_cli"),
+        help="Output folder for PNG renders / artifacts.",
+    )
+    p.add_argument("--tta_lr", type=float, default=1e-6)
+    p.add_argument("--tta_iters", type=int, default=0, help="0 disables TTA.")
+    p.add_argument("--det_thresh", type=float, default=0.7)
+    p.add_argument("--kp_conf_thresh", type=float, default=0.1)
+    p.add_argument("--side_view", action="store_true")
+    p.add_argument("--save_mesh", action="store_true")
+    return p.parse_args()
+
+
+def main() -> int:
+    # Ensure local defaults (GPU if available) but no Space-only preload behavior.
+    os.environ.setdefault("PRIMA_DEMO_MODE", "local")
+    os.environ.setdefault("PRIMA_PRELOAD_ASSETS", "0")
+
+    import numpy as np  # noqa: E402
+
+    import app  # noqa: E402
+
+    args = parse_args()
+    out_dir = Path(args.out)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    img_path = Path(args.image)
+    if not img_path.is_file():
+        raise FileNotFoundError(f"Missing image: {img_path}")
+
+    img_bgr = cv2.imread(str(img_path))
+    if img_bgr is None:
+        raise RuntimeError(f"Failed to read image: {img_path}")
+    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB).astype(np.uint8)
+
+    print("[local_infer] Loading PRIMA model ...")
+    model, model_cfg, renderer, cam_crop_to_full_fn, device = app._load_prima_model()
+    print(f"[local_infer] device={device}")
+
+    print("[local_infer] Building detector (Detectron2 if installed, else SuperAnimal fallback) ...")
+    detector = app._build_detector()
+    print(f"[local_infer] detector={'detectron2' if detector is not None else 'superanimal fallback'}")
+
+    print("[local_infer] Running inference ...")
+    before, after, kpts, mesh_before, mesh_after = app._collect_animal_results(
+        model,
+        model_cfg,
+        renderer,
+        cam_crop_to_full_fn,
+        device,
+        detector,
+        str(out_dir),
+        img_rgb,
+        tta_lr=float(args.tta_lr),
+        tta_num_iters=int(args.tta_iters),
+        det_thresh=float(args.det_thresh),
+        kp_conf_thresh=float(args.kp_conf_thresh),
+        side_view=bool(args.side_view),
+        save_mesh=bool(args.save_mesh),
+    )
+
+    print(f"[local_infer] renders: before={len(before)} after={len(after)} kpts={len(kpts)}")
+    if mesh_before or mesh_after:
+        print(f"[local_infer] meshes: before={mesh_before} after={mesh_after}")
+
+    pngs = sorted(out_dir.glob("*.png"))
+    for p in pngs:
+        print(f"[local_infer] output: {p}")
+
+    if not pngs:
+        raise RuntimeError("No PNG outputs produced.")
+
+    print("[local_infer] OK")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/scripts/run_local_demo_once.py b/scripts/run_local_demo_once.py
new file mode 100644
index 0000000..43b1aa3
--- /dev/null
+++ b/scripts/run_local_demo_once.py
@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+"""
+PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation
+
+Official implementation of the paper:
+"PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation"
+by Xiaohang Yu, Ti Wang, and Mackenzie Weygandt Mathis
+Licensed under a modified MIT license
+
+One-shot local smoke: load PRIMA, run beagle demo (TTA off), print paths to outputs.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+os.environ.setdefault("PRIMA_PRELOAD_ASSETS", "0")
+
+import cv2  # noqa: E402
+
+import app  # noqa: E402
+
+
+def main() -> int:
+    out_dir = ROOT / "demo_out_tta_gradio_local_proof"
+    out_dir.mkdir(parents=True, exist_ok=True)
+    img_path = ROOT / "demo_data" / "beagle.jpg"
+    if not img_path.is_file():
+        print(f"ERROR: missing {img_path}")
+        return 1
+
+    print("[1/4] Loading PRIMA checkpoint …")
+    model, cfg, renderer, device = app._load_prima_model()
+    print(f"      device={device}")
+
+    print("[2/4] Building detector (Detectron2 if installed, else SuperAnimal detector) …")
+    det = app._build_detector()
+    print(f"      detector={'detectron2' if det is not None else 'superanimal fallback'}")
+
+    img = cv2.cvtColor(cv2.imread(str(img_path)), cv2.COLOR_BGR2RGB)
+    print(f"[3/4] Running inference on {img_path.name} (TTA iterations=0) …")
+    before, after, kpts, _, _ = app._collect_animal_results(
+        model,
+        cfg,
+        renderer,
+        device,
+        det,
+        str(out_dir),
+        img,
+        1e-6,
+        0,
+        0.7,
+        0.1,
+        False,
+        False,
+    )
+    print(f"      renders: before={len(before)} after={len(after)} kpts={len(kpts)}")
+
+    pngs = sorted(out_dir.glob("*.png"))
+    print("[4/4] Output files:")
+    for p in pngs:
+        print(f"      {p}")
+
+    if not pngs:
+        print("FAIL: no PNG outputs (often pyrender/display on headless macOS).")
+        return 1
+
+    print("OK: local demo produced outputs.")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/scripts/setup_demo_data.py b/scripts/setup_demo_data.py
new file mode 100644
index 0000000..2ca3659
--- /dev/null
+++ b/scripts/setup_demo_data.py
@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+"""
+PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation
+
+Official implementation of the paper:
+"PRIMA: Boosting Animal Mesh Recovery with Biological Priors and Test-Time Adaptation"
+by Xiaohang Yu, Ti Wang, and Mackenzie Weygandt Mathis
+Licensed under a modified MIT license
+"""
+# Download and arrange PRIMA demo assets into the expected data/ layout.
+# Usage:
+#   python scripts/setup_demo_data.py
+#   python scripts/setup_demo_data.py --include-stage3
+#   python scripts/setup_demo_data.py --force
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parent.parent
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from prima.utils.weights import (
+    DEFAULT_HF_REPO_ID,
+    ensure_demo_assets,
+)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Download PRIMA demo checkpoints and data")
+    parser.add_argument("--data-dir", type=Path, default=Path("data"), help="Target data directory")
+    parser.add_argument("--force", action="store_true", help="Redownload and overwrite existing files")
+    parser.add_argument(
+        "--include-stage3",
+        action="store_true",
+        help="Also prefetch the Stage 3 checkpoint and config",
+    )
+    parser.add_argument(
+        "--hf-repo-id",
+        type=str,
+        default=DEFAULT_HF_REPO_ID,
+        help="Hugging Face repo ID containing demo assets (e.g., org/repo)",
+    )
+    args = parser.parse_args()
+    stages = ("PRIMAS1", "PRIMAS3") if args.include_stage3 else ("PRIMAS1",)
+    ensure_demo_assets(
+        args.data_dir,
+        stages=stages,
+        force=args.force,
+        hf_repo_id=args.hf_repo_id,
+    )
+
+    print("\n[done] Demo assets ready.")
+    print("Run demo:")
+    print("  python demo.py --img_folder demo_data/ --out_folder demo_out/")
+    print("Run demo with TTA:")
+    print("  python demo_tta.py --img_folder demo_data/ --out_folder demo_out_tta/ --tta_lr 1e-6 --tta_num_iters 30")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/scripts/test_local_full.sh b/scripts/test_local_full.sh
new file mode 100755
index 0000000..9bb90fb
--- /dev/null
+++ b/scripts/test_local_full.sh
@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+# Full local smoke: local CLI inference (no Gradio).
+set -euo pipefail
+
+export GIT_TERMINAL_PROMPT=0
+ROOT="$(git rev-parse --show-toplevel)"
+cd "$ROOT"
+VENV="${PRIMA_VENV:-.venv}"
+
+if [[ ! -x "${VENV}/bin/python" ]]; then
+  echo "ERROR: missing ${VENV}. Run: ./scripts/clean_install_local.sh --skip-data" >&2
+  exit 1
+fi
+
+# shellcheck disable=SC1090
+source "${VENV}/bin/activate"
+export PYTHONPATH="${ROOT}${PYTHONPATH:+:${PYTHONPATH}}"
+export PRIMA_PRELOAD_ASSETS=0
+export PRIMA_DEMO_MODE=local
+
+echo "=== [1/3] Demo profile (local) ==="
+python -c "import app; app.get_demo_profile.cache_clear(); p=app.get_demo_profile(); print('profile:', p.mode)"
+
+echo "=== [2/3] DeepLabCut SuperAnimal API ==="
+python -c "
+from deeplabcut.pose_estimation_pytorch.apis import superanimal_analyze_images  # noqa: F401
+print('DeepLabCut SuperAnimal: OK')
+"
+
+echo "=== [3/3] PRIMA local CLI inference (beagle, TTA off) ==="
+python "${ROOT}/scripts/local_infer.py" --tta_iters 0 2>&1
+
+echo "=== All local checks passed ==="
diff --git a/scripts/update_headers.py b/scripts/update_headers.py
index ab35d63..96b8b16 100644
--- a/scripts/update_headers.py
+++ b/scripts/update_headers.py
@@ -40,24 +40,13 @@ def should_skip_file(file_path):
     Returns:
         True if the file should be skipped, False otherwise
     """
-    skip_dirs = {'.git', '__pycache__', '.pytest_cache', 'venv', 'env', '.tox', 'build', 'dist', '.eggs'}
+    skip_dirs = {'.git', '__pycache__', '.pytest_cache', 'venv', '.venv', 'env', '.tox', 'build', 'dist', '.eggs', 'site-packages'}
     
     # Skip if in excluded directory
     for part in file_path.parts:
         if part in skip_dirs:
             return True
     
-    # Skip __init__.py files that are typically minimal
-    if file_path.name == '__init__.py':
-        try:
-            with open(file_path, 'r', encoding='utf-8') as f:
-                content = f.read()
-            # Skip if __init__.py is very short (likely just imports)
-            if len(content.strip()) < 50:
-                return True
-        except Exception:
-            pass
-    
     return False
 
 
@@ -75,25 +64,13 @@ def has_header(content):
     if STANDARD_HEADER.strip() in content:
         return True
     
-    # Check for header with additional content (like in sort.py)
-    # Header should contain the key elements
-    lines = content.split('\n')
-    if len(lines) < 3:
-        return False
-    
-    # Check if it starts with a docstring
-    if not lines[0].strip().startswith('"""'):
-        return False
-    
-    # Check for key header components in the first 15 lines
-    header_section = '\n'.join(lines[:15])
     required_elements = [
-        'FMPose3D: monocular 3D Pose Estimation via Flow Matching',
-        'Ti Wang, Xiaohang Yu, and Mackenzie Weygandt Mathis',
-        'Licensed under Apache 2.0'
+        'PRIMA: Boosting Animal Mesh Recovery with Biological Priors',
+        'Xiaohang Yu, Ti Wang, and Mackenzie Weygandt Mathis',
+        'Licensed under a modified MIT license',
     ]
-    
-    return all(elem in header_section for elem in required_elements)
+
+    return all(elem in content for elem in required_elements)
 
 
 def needs_header_update(content):
@@ -167,11 +144,12 @@ def add_or_update_header(file_path, check_only=False):
                         break
                 
                 if future_import_index is not None:
-                    # If there's a from __future__ import, add header AFTER it
-                    new_lines.extend(lines[insert_index:future_import_index+1])
-                    new_lines.append(STANDARD_HEADER)
-                    new_lines.append('')
-                    new_lines.extend(lines[future_import_index+1:])
+                    # ``from __future__`` must stay first; PRIMA header docstring follows it.
+                    new_lines.extend(lines[insert_index:future_import_index + 1])
+                    if STANDARD_HEADER.strip() not in content:
+                        new_lines.append(STANDARD_HEADER)
+                        new_lines.append('')
+                    new_lines.extend(lines[future_import_index + 1:])
                 else:
                     # Otherwise, add header at the beginning (after shebang if present)
                     new_lines.append(STANDARD_HEADER)