Add ZeroGPU decorator for GPU function detection

app.py

@@ -83,72 +83,89 @@ def keep_largest_component(mask):
         return mask
 
 
-def segment_slice_api(
-    nifti_b64: str,
-    slice_idx: int,
-    text_prompt: str = "tumor",
-    checkpoint: str = "brain",
-):
-    """
-    Segment a single slice from a NIfTI volume.
-    Currently uses fallback segmentation (SAM3 to be integrated).
-    """
-    start_time = time.time()
-
-    try:
-        import nibabel as nib
-
-        # Decode NIfTI
-        nifti_bytes = base64.b64decode(nifti_b64)
-
-        with tempfile.NamedTemporaryFile(suffix=".nii.gz", delete=False) as f:
-            f.write(nifti_bytes)
-            temp_path = f.name
-
-        nii = nib.load(temp_path)
-        volume = nii.get_fdata().astype(np.float32)
-        os.unlink(temp_path)
+def _do_segment(nifti_b64, slice_idx, text_prompt, checkpoint):
+    """Core segmentation logic."""
+    import nibabel as nib
 
-        logger.info(
-            f"Loaded volume shape: {volume.shape}, segmenting slice {slice_idx}"
-        )
-
-        # Validate slice index
-        if slice_idx < 0 or slice_idx >= volume.shape[0]:
-            return {
-                "success": False,
-                "error": f"Slice index {slice_idx} out of range [0, {volume.shape[0]})",
-            }
+    start_time = time.time()
 
-        # Extract slice and segment
-        slice_2d = volume[slice_idx]
-        mask = fallback_segment(slice_2d)
-        mask = keep_largest_component(mask)
+    # Decode NIfTI
+    nifti_bytes = base64.b64decode(nifti_b64)
 
-        # Extract contours
-        contours = find_contours(mask)
+    with tempfile.NamedTemporaryFile(suffix=".nii.gz", delete=False) as f:
+        f.write(nifti_bytes)
+        temp_path = f.name
 
-        # Encode mask as base64
-        mask_b64 = base64.b64encode(mask.tobytes()).decode()
+    nii = nib.load(temp_path)
+    volume = nii.get_fdata().astype(np.float32)
+    os.unlink(temp_path)
 
-        inference_time = int((time.time() - start_time) * 1000)
-        logger.info(
-            f"Segmented slice {slice_idx} in {inference_time}ms, mask sum: {mask.sum()}"
-        )
+    logger.info(f"Loaded volume shape: {volume.shape}, segmenting slice {slice_idx}")
 
+    # Validate slice index
+    if slice_idx < 0 or slice_idx >= volume.shape[0]:
         return {
-            "success": True,
-            "backend": "fallback",
-            "mask_b64": mask_b64,
-            "mask_shape": list(mask.shape),
-            "contours": contours,
-            "slice_idx": slice_idx,
-            "inference_time_ms": inference_time,
+            "success": False,
+            "error": f"Slice index {slice_idx} out of range [0, {volume.shape[0]})",
         }
 
-    except Exception as e:
-        logger.error(f"Segmentation failed: {e}")
-        return {"success": False, "error": str(e)}
+    # Extract slice and segment
+    slice_2d = volume[slice_idx]
+    mask = fallback_segment(slice_2d)
+    mask = keep_largest_component(mask)
+
+    # Extract contours
+    contours = find_contours(mask)
+
+    # Encode mask as base64
+    mask_b64 = base64.b64encode(mask.tobytes()).decode()
+
+    inference_time = int((time.time() - start_time) * 1000)
+    logger.info(
+        f"Segmented slice {slice_idx} in {inference_time}ms, mask sum: {mask.sum()}"
+    )
+
+    return {
+        "success": True,
+        "backend": "fallback",
+        "mask_b64": mask_b64,
+        "mask_shape": list(mask.shape),
+        "contours": contours,
+        "slice_idx": slice_idx,
+        "inference_time_ms": inference_time,
+    }
+
+
+# Create GPU-decorated and non-GPU versions of the API
+if ZEROGPU_AVAILABLE:
+
+    @spaces.GPU(duration=30)
+    def segment_slice_api(
+        nifti_b64: str,
+        slice_idx: int,
+        text_prompt: str = "tumor",
+        checkpoint: str = "brain",
+    ):
+        """Segment a single slice (GPU-accelerated when available)."""
+        try:
+            return _do_segment(nifti_b64, slice_idx, text_prompt, checkpoint)
+        except Exception as e:
+            logger.error(f"Segmentation failed: {e}")
+            return {"success": False, "error": str(e)}
+else:
+
+    def segment_slice_api(
+        nifti_b64: str,
+        slice_idx: int,
+        text_prompt: str = "tumor",
+        checkpoint: str = "brain",
+    ):
+        """Segment a single slice."""
+        try:
+            return _do_segment(nifti_b64, slice_idx, text_prompt, checkpoint)
+        except Exception as e:
+            logger.error(f"Segmentation failed: {e}")
+            return {"success": False, "error": str(e)}
 
 
 def load_and_display_nifti(file):