0

app.py

@@ -11,7 +11,7 @@ import tempfile
 import threading
 import uuid
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Sequence, Tuple
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
 
 import cv2
 import gradio as gr
@@ -50,6 +50,12 @@ GRANULARITY_MIN = float(os.getenv("UNSAMV2_GRAN_MIN", 0.1))
 GRANULARITY_MAX = float(os.getenv("UNSAMV2_GRAN_MAX", 1.0))
 ZERO_GPU_ENABLED = os.getenv("UNSAMV2_ENABLE_ZEROGPU", "1").lower() in {"1", "true", "yes"}
 ZERO_GPU_DURATION = int(os.getenv("UNSAMV2_ZEROGPU_DURATION", "60"))
+ZERO_GPU_WHOLE_DURATION = int(
+    os.getenv("UNSAMV2_ZEROGPU_WHOLE_DURATION", str(ZERO_GPU_DURATION))
+)
+ZERO_GPU_VIDEO_DURATION = int(
+    os.getenv("UNSAMV2_ZEROGPU_VIDEO_DURATION", str(max(120, ZERO_GPU_DURATION)))
+)
 MAX_VIDEO_FRAMES = int(os.getenv("UNSAMV2_MAX_VIDEO_FRAMES", "360"))
 WHOLE_IMAGE_POINTS_PER_SIDE = int(os.getenv("UNSAMV2_WHOLE_POINTS", "64"))
 WHOLE_IMAGE_MAX_MASKS = 1000
@@ -360,6 +366,20 @@ def choose_device() -> torch.device:
     return torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
+def wrap_with_zero_gpu(
+    fn: Callable[..., Any],
+    duration: int,
+) -> Callable[..., Any]:
+    if spaces is None or not ZERO_GPU_ENABLED:
+        return fn
+    try:
+        LOGGER.info("Enabling ZeroGPU (duration=%ss) for %s", duration, fn.__name__)
+        return spaces.GPU(duration=duration)(fn)  # type: ignore[misc]
+    except Exception:  # pragma: no cover - defensive logging
+        LOGGER.exception("Failed to wrap %s with ZeroGPU; running on CPU", fn.__name__)
+        return fn
+
+
 def build_granularity_tensor(value: float, device: torch.device) -> torch.Tensor:
     tensor = torch.tensor([[[[value]]]], dtype=torch.float32, device=device)
     return tensor
@@ -835,10 +855,19 @@ def run_video_frame_segmentation(
     return overlay, status
 
 
-if spaces is not None and ZERO_GPU_ENABLED:
-    segment_fn = spaces.GPU(duration=ZERO_GPU_DURATION)(_run_segmentation)
-else:
-    segment_fn = _run_segmentation
+segment_fn = wrap_with_zero_gpu(_run_segmentation, ZERO_GPU_DURATION)
+whole_image_fn = wrap_with_zero_gpu(
+    run_whole_image_segmentation,
+    ZERO_GPU_WHOLE_DURATION,
+)
+video_frame_fn = wrap_with_zero_gpu(
+    run_video_frame_segmentation,
+    ZERO_GPU_VIDEO_DURATION,
+)
+video_segmentation_fn = wrap_with_zero_gpu(
+    run_video_segmentation,
+    ZERO_GPU_VIDEO_DURATION,
+)
 
 
 def build_demo() -> gr.Blocks:
@@ -1018,7 +1047,7 @@ def build_demo() -> gr.Blocks:
                 whole_status = gr.Markdown(" Ready for whole-image masks.")
 
                 whole_generate_btn.click(
-                    run_whole_image_segmentation,
+                    whole_image_fn,
                     inputs=[
                         whole_image_input,
                         whole_granularity,
@@ -1162,7 +1191,7 @@ def build_demo() -> gr.Blocks:
                 )
 
                 video_frame_btn.click(
-                    run_video_frame_segmentation,
+                    video_frame_fn,
                     inputs=[
                         video_state,
                         video_points_state,
@@ -1201,7 +1230,7 @@ def build_demo() -> gr.Blocks:
                 )
 
                 video_segment_btn.click(
-                    run_video_segmentation,
+                    video_segmentation_fn,
                     inputs=[
                         video_state,
                         video_points_state,