startup.py

"""
startup.py — automatic model bootstrap for HuggingFace Spaces (and local dev).

Download strategy for MedSAM checkpoint:
  1. MEDSAM_CHECKPOINT env var set and file exists   → use it as-is (local / server)
  2. HuggingFace Hub: wanglab/medsam-vit-base        → hf_hub_download (same as MedGemma)
  3. Zenodo fallback                                  → urllib download, cached locally

After downloading, os.environ["MEDSAM_CHECKPOINT"] is updated so the rest of the
app always sees the correct resolved path.
"""
from __future__ import annotations

import os
import urllib.request
from pathlib import Path

# ---------------------------------------------------------------------------
# Config (all overridable via env vars)
# ---------------------------------------------------------------------------

_MEDSAM_HF_REPO      = os.getenv("MEDSAM_HF_REPO", "wanglab/medsam-vit-base")
_MEDSAM_HF_FILENAME  = os.getenv("MEDSAM_HF_FILENAME", "medsam_vit_b.pth")
_MEDSAM_ZENODO_URL   = (
    "https://zenodo.org/records/10689643/files/medsam_vit_b.pth"
)

# Cache inside HF_HOME so persistent storage on HF Spaces is respected
_CACHE_DIR = (
    Path(os.getenv("HF_HOME", str(Path.home() / ".cache" / "huggingface")))
    / "medsam"
)


# ---------------------------------------------------------------------------
# MedSAM checkpoint resolution
# ---------------------------------------------------------------------------

def _resolve_medsam_checkpoint() -> str:
    """Return the local path to medsam_vit_b.pth, downloading if needed."""

    # ── 1. Explicit path already provided and file exists ────────────────────
    explicit = os.getenv("MEDSAM_CHECKPOINT", "")
    if explicit and Path(explicit).exists():
        print(f"[startup] MedSAM checkpoint: using MEDSAM_CHECKPOINT={explicit}")
        return explicit

    _CACHE_DIR.mkdir(parents=True, exist_ok=True)

    # ── 2. HuggingFace Hub — wanglab/medsam-vit-base (same mechanism as MedGemma) ──
    try:
        from huggingface_hub import hf_hub_download
        print(
            f"[startup] Downloading MedSAM from HuggingFace Hub "
            f"({_MEDSAM_HF_REPO}/{_MEDSAM_HF_FILENAME})…"
        )
        path = hf_hub_download(
            repo_id=_MEDSAM_HF_REPO,
            filename=_MEDSAM_HF_FILENAME,
        )
        print(f"[startup] MedSAM cached at: {path}")
        os.environ["MEDSAM_CHECKPOINT"] = path
        return path
    except Exception as hf_err:
        print(
            f"[startup] HuggingFace download failed ({hf_err}), "
            f"falling back to Zenodo…"
        )

    # ── 3. Zenodo fallback with local cache ───────────────────────────────────
    cached = _CACHE_DIR / _MEDSAM_HF_FILENAME
    if cached.exists():
        print(f"[startup] MedSAM checkpoint found in cache: {cached}")
        os.environ["MEDSAM_CHECKPOINT"] = str(cached)
        return str(cached)

    print(f"[startup] Downloading MedSAM checkpoint (~375 MB) from Zenodo…")

    def _progress(count: int, block: int, total: int) -> None:
        if total > 0 and count % 500 == 0:
            pct = min(100, count * block * 100 // total)
            print(f"\r[startup]   {pct}%", end="", flush=True)

    urllib.request.urlretrieve(_MEDSAM_ZENODO_URL, cached, reporthook=_progress)
    print(f"\n[startup] Saved to: {cached}")
    os.environ["MEDSAM_CHECKPOINT"] = str(cached)
    return str(cached)


# ---------------------------------------------------------------------------
# Public entry-point
# ---------------------------------------------------------------------------

def initialize_all_models(store: dict) -> str:
    """
    Load MedGemma and MedSAM into *store*.
    Returns an HTML string for the load_status component.
    """
    from src.config.endpoints import MEDGEMMA_MODEL_ID, MEDSAM_DEVICE

    lines: list[str] = []

    # ── MedGemma (Detection + Report generation) ──────────────────────────────
    try:
        from src.clients.medgemma_client import MedGemmaClient
        print("[startup] Loading MedGemma…")
        store["medgemma"] = MedGemmaClient(model_id=MEDGEMMA_MODEL_ID)
        lines.append(
            "<span style='color:#2E7D32;font-weight:600'>✅ Detection model ready</span>"
        )
        print("[startup] MedGemma ready.")
    except Exception as exc:
        lines.append(
            f"<span style='color:#C62828;font-weight:600'>❌ Detection model failed: {exc}</span>"
        )
        print(f"[startup] MedGemma failed: {exc}")

    # ── MedSAM (Segmentation) ─────────────────────────────────────────────────
    try:
        ckpt_path = _resolve_medsam_checkpoint()
        store["medsam_ckpt_path"] = ckpt_path          # expose for lazy-load fallback

        from src.clients.medsam_client import MedSAMClient
        print(f"[startup] Loading MedSAM from {ckpt_path}…")
        store["medsam"] = MedSAMClient(checkpoint_path=ckpt_path, device=MEDSAM_DEVICE)
        lines.append(
            f"<span style='color:#2E7D32;font-weight:600'>"
            f"✅ Segmentation model ready ({MEDSAM_DEVICE})</span>"
        )
        print("[startup] MedSAM ready.")
    except Exception as exc:
        lines.append(
            f"<span style='color:#C62828;font-weight:600'>"
            f"❌ Segmentation model failed: {exc}</span>"
        )
        print(f"[startup] MedSAM failed: {exc}")

    return "<br>".join(lines)