formscout/startup.py

"""
Checkpoint bootstrap — downloads missing model files from HF model repo on first run.
Called once at app startup before build_app(); no-ops if files already present.
"""
from __future__ import annotations

import logging
from pathlib import Path

logger = logging.getLogger(__name__)

CHECKPOINT_REPO = "silas-therapy/formscout-checkpoints"
ROOT = Path(__file__).parent.parent

_CHECKPOINTS = [
    "checkpoints/yolo26/yolo26n-pose.pt",
    "checkpoints/yolo26/yolo26s-pose.pt",
    "checkpoints/yolo26/yolo26m-pose.pt",
    "checkpoints/yolo26/yolo26l-pose.pt",
    "checkpoints/yolo26/yolo26x-pose.pt",
    "checkpoints/mediapipe/pose_landmarker_full.task",
]


def ensure_checkpoints() -> None:
    """Download any missing checkpoints from silas-therapy/formscout-checkpoints."""
    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        logger.warning("huggingface_hub not installed — skipping checkpoint download")
        return

    for rel_path in _CHECKPOINTS:
        local = ROOT / rel_path
        if local.exists():
            continue
        logger.info("Downloading %s ...", rel_path)
        try:
            local.parent.mkdir(parents=True, exist_ok=True)
            hf_hub_download(
                repo_id=CHECKPOINT_REPO,
                filename=rel_path,
                local_dir=str(ROOT),
            )
            logger.info("Downloaded %s", rel_path)
        except Exception as e:
            logger.warning("Failed to download %s: %s", rel_path, e)