Update tools/startup_selfcheck.py

tools/startup_selfcheck.py

@@ -1,54 +1,88 @@
 #!/usr/bin/env python3
+# tools/startup_selfcheck.py
 """
-Startup Self-Check for Hugging Face Space
-- Verifies SAM2 + MatAnyone load and run on a dummy frame
-- Prints ✅ / ❌ in logs
+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
 
-from models.loaders.sam2_loader import SAM2Loader
-from models.loaders.matanyone_loader import MatAnyoneLoader
+# 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
 
-logger = logging.getLogger("selfcheck")
-if not logger.handlers:
+log = logging.getLogger("selfcheck")
+if not log.handlers:
     logging.basicConfig(level=logging.INFO)
 
-def run_selfcheck():
-    results = {}
+def run_selfcheck() -> None:
+    """Do the actual check; log results. Non-throwing."""
     try:
-        # tiny dummy image
+        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 ---
-        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)
-            results["SAM2"] = out["masks"].shape
-            logger.info(f"✅ SAM2 OK, masks shape {out['masks'].shape}")
-        else:
-            results["SAM2"] = None
-            logger.error("❌ SAM2 failed to load")
+        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 ---
-        mat_loader = MatAnyoneLoader(device="cuda" if torch.cuda.is_available() else "cpu")
-        session = mat_loader.load()
-        if session:
-            alpha = session(img, np.ones((64, 64), np.float32))  # coarse mask
-            results["MatAnyone"] = alpha.shape
-            logger.info(f"✅ MatAnyone OK, alpha shape {alpha.shape}")
-        else:
-            results["MatAnyone"] = None
-            logger.error("❌ MatAnyone failed to load")
+        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:
-        logger.error(f"❌ Self-check error: {e}", exc_info=True)
+        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
 
-    return results
+    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")
 
-if __name__ == "__main__":
-    run_selfcheck()
+    t = threading.Thread(target=_worker, name="startup-selfcheck", daemon=True)
+    t.start()