Cap 5B ORT cache + 6 more demo images

app.py

@@ -233,6 +233,12 @@ def _get_dense_model(task: str, size: str):
     os.makedirs(local_dir, exist_ok=True)
     ckpt = hf_hub_download(repo_id=spec["repo"], filename=spec["filename"], local_dir=local_dir)
 
+    # If a 5B ORT session (5-6 GB) is resident, evict it before init_model allocates the 1B dense weights.
+    if size == "1b" and _ORT_SESSIONS:
+        _ORT_SESSIONS.clear()
+        import gc
+        gc.collect()
+
     model = init_model(config, ckpt, device="cpu")
 
     while len(_MODELS) >= _MAX_CACHED:
@@ -273,6 +279,13 @@ def _get_pose_model(size: str):
     local_dir = f"/tmp/sapiens_models/pose-{size}"
     os.makedirs(local_dir, exist_ok=True)
     ckpt = hf_hub_download(repo_id=spec["repo"], filename=spec["filename"], local_dir=local_dir)
+
+    # Drop a resident 5B ORT session before allocating the 1B pose weights.
+    if size == "1b" and _ORT_SESSIONS:
+        _ORT_SESSIONS.clear()
+        import gc
+        gc.collect()
+
     model = init_model(config, ckpt, device="cpu")
 
     codec_cfg = dict(model.cfg.codec)
@@ -396,7 +409,9 @@ def _infer_pointmap(image_bgr, model):
 
 # --- 5B INT8 ONNX path -------------------------------------------------------
 def _get_ort_session(task: str):
-    """Lazy-load + cache an ORT session for {task}_5b_int8.onnx."""
+    """Lazy-load + cache an ORT session for {task}_5b_int8.onnx.
+    Each 5B session is 5-6 GB RAM. cpu-basic has 16 GB total, so keep at most one
+    5B session live and evict cached dense/pose PyTorch models that would push us OOM."""
     key = (task, "5b")
     sess = _ORT_SESSIONS.get(key)
     if sess is not None:
@@ -406,10 +421,21 @@ def _get_ort_session(task: str):
     spec = VARIANTS[key]
     cache_dir = os.environ.get("ONNX_5B_CACHE", "/app/onnx_5b")
     os.makedirs(cache_dir, exist_ok=True)
-    # Download both the graph .onnx and its external-data sidecar.
     fn = spec["onnx_filename"]
     onnx_path = hf_hub_download(repo_id=spec["onnx_repo"], filename=fn, local_dir=cache_dir)
     hf_hub_download(repo_id=spec["onnx_repo"], filename=fn + ".data", local_dir=cache_dir)
+    # Evict any prior 5B ORT session and any 1b dense models — they together exceed 16 GB.
+    import gc
+    if _ORT_SESSIONS:
+        _ORT_SESSIONS.clear()
+        gc.collect()
+    for k in list(_MODELS.keys()):
+        if k[1] in ("1b", "0.8b"):
+            del _MODELS[k]
+    for k in list(_POSE_MODELS.keys()):
+        if k[1] in ("1b", "0.8b"):
+            del _POSE_MODELS[k]
+    gc.collect()
     sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
     _ORT_SESSIONS[key] = sess
     return sess