Bind ZeroGPU handler directly and log SAM stages

app.py

@@ -233,10 +233,17 @@ def _load_models_at_startup() -> None:
 _load_models_at_startup()
 
 
-def _segment(image: Image.Image, prompt: str) -> tuple[np.ndarray | None, int]:
+def _segment(
+    image: Image.Image, prompt: str, request_id: str
+) -> tuple[np.ndarray | None, int, float | None]:
     import torch
 
     model, processor = _load_sam3()
+    _request_log(
+        request_id,
+        "SAM 3 inference started",
+        f"prompt={prompt!r}; image={image.width}x{image.height}",
+    )
     inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
     with torch.inference_mode():
         output = model(**inputs)
@@ -248,10 +255,17 @@ def _segment(image: Image.Image, prompt: str) -> tuple[np.ndarray | None, int]:
     )[0]
     masks = result["masks"]
     if len(masks) == 0:
-        return None, 0
+        _request_log(request_id, "SAM 3 inference completed", "instances=0")
+        return None, 0, None
     scores = torch.as_tensor(result["scores"])
     best = int(scores.argmax().item())
-    return masks[best].detach().cpu().numpy().astype(bool), len(masks)
+    best_score = float(scores[best].item())
+    _request_log(
+        request_id,
+        "SAM 3 inference completed",
+        f"instances={len(masks)}; selected_mask={best}; score={best_score:.4f}",
+    )
+    return masks[best].detach().cpu().numpy().astype(bool), len(masks), best_score
 
 
 def _overlay(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
@@ -282,9 +296,28 @@ def _request_log(request_id: str, event: str, detail: str = "") -> None:
 
 
 @spaces.GPU(duration=360)
-def _reconstruct_gpu(
-    image: Image.Image, prompt: str, request_id: str
+def reconstruct(
+    image: Image.Image | None, prompt: str
 ) -> tuple[Any, Any, Any, str]:
+    """Gradio's direct ZeroGPU event handler; keep it directly decorated."""
+    if image is None:
+        return None, None, None, "Upload an image first."
+    prompt = (prompt or "").strip()
+    if not prompt:
+        return None, None, None, "Enter a short object description, such as 'chair'."
+    if SETUP_ERROR:
+        return None, None, None, f"Space setup is incomplete: {SETUP_ERROR}"
+    if MODEL_SETUP_ERROR:
+        return None, None, None, f"Model setup failed: {MODEL_SETUP_ERROR}"
+    if SAM3 is None or SAM3D is None:
+        return None, None, None, "Models are not ready; check the Space startup logs."
+
+    request_id = uuid.uuid4().hex[:8]
+    _request_log(
+        request_id,
+        "GPU handler entered",
+        "maximum_gpu_seconds=360",
+    )
     started = time.monotonic()
     rgb = image.convert("RGB")
     image_array = np.asarray(rgb, dtype=np.uint8)
@@ -297,11 +330,11 @@ def _reconstruct_gpu(
     try:
         _request_log(request_id, "stage started", stage)
         stage_started = time.monotonic()
-        mask, count = _segment(rgb, prompt)
+        mask, count, score = _segment(rgb, prompt, request_id)
         _request_log(
             request_id,
-            "stage completed",
-            f"{stage}; matching_instances={count}; stage_seconds={time.monotonic() - stage_started:.1f}; total_seconds={time.monotonic() - started:.1f}",
+            "SAM 3 stage completed",
+            f"matching_instances={count}; selected_score={score}; stage_seconds={time.monotonic() - stage_started:.1f}; total_seconds={time.monotonic() - started:.1f}",
         )
         if mask is None:
             message = f"SAM 3 did not find '{prompt}' in this image."
@@ -356,64 +389,6 @@ def _reconstruct_gpu(
         )
 
 
-def reconstruct(image: Image.Image | None, prompt: str):
-    """Stream a start message, then call the bounded GPU function."""
-    if image is None:
-        yield None, None, None, "Upload an image first."
-        return
-    prompt = (prompt or "").strip()
-    if not prompt:
-        yield None, None, None, "Enter a short object description, such as 'chair'."
-        return
-    if SETUP_ERROR:
-        yield None, None, None, f"Space setup is incomplete: {SETUP_ERROR}"
-        return
-    if MODEL_SETUP_ERROR:
-        yield None, None, None, f"Model setup failed: {MODEL_SETUP_ERROR}"
-        return
-    if SAM3 is None or SAM3D is None:
-        yield None, None, None, "Models are not ready; check the Space startup logs."
-        return
-
-    request_id = uuid.uuid4().hex[:8]
-    _request_log(
-        request_id,
-        "request accepted",
-        "submitting to ZeroGPU; maximum GPU time is 360s",
-    )
-    yield None, None, None, "Queued on ZeroGPU; maximum GPU time is 360 seconds. Progress is recorded in Space Container logs."
-    try:
-        result = _reconstruct_gpu(image, prompt, request_id)
-    except Exception as exc:
-        _request_log(
-            request_id,
-            "ZeroGPU call failed",
-            f"{type(exc).__name__}: {exc}",
-        )
-        traceback.print_exc()
-        error_text = str(exc)
-        error_lower = error_text.lower()
-        is_timeout = isinstance(exc, TimeoutError) or any(
-            word in error_lower for word in ("timeout", "timed out", "deadline", "time limit")
-        )
-        is_abort = any(word in error_lower for word in ("aborted", "killed", "worker lost"))
-        if is_timeout:
-            status_text = (
-                "ZeroGPU reported that the request hit its runtime limit (maximum 360 seconds). "
-                "Check Container logs for the last completed stage."
-            )
-        elif is_abort:
-            status_text = (
-                "ZeroGPU aborted or lost the request before it returned; that alone does not prove "
-                "a timeout. Check Container logs for the last completed stage."
-            )
-        else:
-            status_text = f"ZeroGPU request failed: {type(exc).__name__}: {error_text}"
-        yield None, None, None, status_text
-        return
-    yield result
-
-
 def _readiness() -> str:
     if SETUP_ERROR:
         return f"Setup: {SETUP_ERROR}"