Add step-by-step ZeroGPU diagnostics

app.py

@@ -124,6 +124,75 @@ def _embed(masks, img_rgb, df_base):
         raise
 
 
+# ── ZeroGPU diagnostics ──────────────────────────────────────────────────────
+# "GPU task aborted" means the worker process died, so nothing it printed
+# survives. Each step below is its own @spaces.GPU call: whichever one fails to
+# come back is the culprit.
+
+@spaces.GPU(duration=60)
+def _diag_alloc():
+    import torch
+    x = torch.zeros(2048, 2048, device="cuda")
+    return f"{torch.cuda.get_device_name(0)}, alloc ok, sum={float(x.sum())}"
+
+
+@spaces.GPU(duration=60)
+def _diag_devices():
+    import torch
+    out = []
+    try:
+        p = next(CELLPOSE_MODEL.net.parameters())
+        out.append(f"cellpose net on {p.device} ({p.dtype})")
+    except Exception as e:
+        out.append(f"cellpose net ERR {type(e).__name__}: {e}")
+    try:
+        p = next(DINOV2_MODEL.parameters())
+        out.append(f"dinov2 on {p.device}")
+    except Exception as e:
+        out.append(f"dinov2 ERR {type(e).__name__}: {e}")
+    out.append(f"torch {torch.__version__}")
+    return "; ".join(out)
+
+
+@spaces.GPU(duration=120)
+def _diag_cellpose():
+    import numpy as np
+    img = (np.random.rand(256, 256, 3) * 255).astype(np.uint8)
+    m = rs.stage_segment(img, gpu=True, cellpose_model=CELLPOSE_MODEL)
+    return f"segmented 256x256 noise -> {int(m.max())} labels, dtype {m.dtype}"
+
+
+@spaces.GPU(duration=120)
+def _diag_dino():
+    import numpy as np
+    from rootscope.cnn_embeddings import extract_cnn_embeddings
+    masks = np.zeros((128, 128), dtype=np.int32)
+    masks[20:60, 20:60] = 1
+    masks[70:110, 70:110] = 2
+    img = (np.random.rand(128, 128, 3) * 255).astype(np.uint8)
+    df = extract_cnn_embeddings(masks, img, use_gpu=True, model=DINOV2_MODEL)
+    return f"embedded {len(df)} cells, {df.shape[1] - 1} dims"
+
+
+def run_diagnostics():
+    steps = [
+        ("1. CUDA allocation", _diag_alloc),
+        ("2. model devices", _diag_devices),
+        ("3. Cellpose-SAM on GPU", _diag_cellpose),
+        ("4. DINOv2 on GPU", _diag_dino),
+    ]
+    lines = []
+    for name, fn in steps:
+        try:
+            lines.append(f"PASS  {name}: {fn()}")
+        except Exception as e:
+            lines.append(f"FAIL  {name}: {type(e).__name__}: {e}")
+            lines.append("")
+            lines.append("^ this step is where it breaks.")
+            break
+    return "\n".join(lines)
+
+
 # ── helpers ──────────────────────────────────────────────────────────────────
 
 def read_um_per_px(tif_path):
@@ -379,6 +448,10 @@ with gr.Blocks(title="RootScope", **_blocks_kwargs) as demo:
                          "until predictions stop changing.",
                 )
             run_btn = gr.Button("Run RootScope", variant="primary")
+            with gr.Accordion("Diagnostics", open=False):
+                diag_btn = gr.Button("Test the GPU step by step")
+                diag_out = gr.Textbox(label="Result", lines=7,
+                                      show_copy_button=True)
 
         with gr.Column(scale=1):
             overlay = gr.Image(label="Predicted cell types", type="filepath",
@@ -433,6 +506,7 @@ with gr.Blocks(title="RootScope", **_blocks_kwargs) as demo:
     )
     model_pick.change(switch_model, [results_state, model_pick],
                       [overlay, summary, stats_md, csv_btn])
+    diag_btn.click(run_diagnostics, None, diag_out)
 
 if __name__ == "__main__":
     launch_kwargs = {"theme": _THEME} if _THEME_ON_LAUNCH else {}