tools/startup_selfcheck.py

#!/usr/bin/env python3
# tools/startup_selfcheck.py
"""
Quiet Startup Self-Check for Hugging Face Space
- Verifies SAM2 + MatAnyone can load and run a tiny inference
- Defaults to ASYNC so UI launches immediately; logs results
- Switch to SYNC by setting env SELF_CHECK_MODE=sync
"""

import logging
import os
import threading
import numpy as np
import torch

# Keep imports lazy so a busted loader doesn't break the Space
def _safe_imports():
    from models.loaders.sam2_loader import SAM2Loader
    from models.loaders.matanyone_loader import MatAnyoneLoader
    return SAM2Loader, MatAnyoneLoader

log = logging.getLogger("selfcheck")
if not log.handlers:
    logging.basicConfig(level=logging.INFO)

def run_selfcheck() -> None:
    """Do the actual check; log results. Non-throwing."""
    try:
        SAM2Loader, MatAnyoneLoader = _safe_imports()
    except Exception as e:
        log.error(f"❌ Self-check import failure: {e}", exc_info=True)
        return

    try:
        # tiny dummy frame + mask
        img = np.zeros((64, 64, 3), dtype=np.uint8)
        mask = np.ones((64, 64), dtype=np.float32)

        # --- SAM2 ---
        try:
            sam_loader = SAM2Loader(device="cuda" if torch.cuda.is_available() else "cpu")
            sam = sam_loader.load("tiny")
            if sam:
                sam.set_image(img)
                out = sam.predict(point_coords=None, point_labels=None)
                m = out["masks"]
                log.info(f"✅ SAM2 OK — masks {m.shape} {m.dtype} [{m.min():.3f},{m.max():.3f}]")
            else:
                log.error("❌ SAM2 failed to load (adapter is None)")
        except Exception as e:
            log.error(f"❌ SAM2 self-check error: {e}", exc_info=True)

        # --- MatAnyone ---
        try:
            mat_loader = MatAnyoneLoader(device="cuda" if torch.cuda.is_available() else "cpu")
            session = mat_loader.load()
            if session:
                alpha = session(img, mask)  # first frame needs a mask (fallback ok)
                log.info(f"✅ MatAnyone OK — alpha {alpha.shape} {alpha.dtype} [{alpha.min():.3f},{alpha.max():.3f}]")
            else:
                log.error("❌ MatAnyone failed to load (session is None)")
        except Exception as e:
            log.error(f"❌ MatAnyone self-check error: {e}", exc_info=True)

    except Exception as e:
        log.error(f"❌ Self-check unexpected error: {e}", exc_info=True)

def schedule_startup_selfcheck(mode: str = "async") -> None:
    """
    mode: "async" (default) → run in a daemon thread; UI doesn't wait
          "sync"           → block until check completes (slower startup)
    """
    mode = (mode or "async").lower()
    if mode == "sync":
        log.info("🔎 Running model self-check (sync)…")
        run_selfcheck()
        return

    def _worker():
        try:
            log.info("🔎 Running model self-check (async)…")
            run_selfcheck()
        except Exception:
            # Make absolutely sure this never crashes the app
            log.exception("Self-check thread crashed")

    t = threading.Thread(target=_worker, name="startup-selfcheck", daemon=True)
    t.start()