Upload folder using huggingface_hub

app.py

@@ -34,6 +34,15 @@ from autolabel.export import build_coco
 from autolabel.segment import load_sam2, segment_with_boxes
 from autolabel.utils import save_json, setup_logging
 
+# HuggingFace ZeroGPU: GPU is only available inside @spaces.GPU-decorated
+# functions.  Import gracefully so the app still runs locally without it.
+try:
+    import spaces as _spaces  # type: ignore
+    _ZERO_GPU = True
+except ImportError:
+    _spaces = None
+    _ZERO_GPU = False
+
 setup_logging(logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -162,32 +171,70 @@ def _get_sam2():
 # Shared inference helpers
 # ---------------------------------------------------------------------------
 
-def _run_detection(
+def _infer_on_device(
     pil_image: Image.Image,
     prompts: list[str],
     threshold: float,
     mode: str,
+    device: str,
+    dtype,
 ) -> list[dict]:
-    """Run OWLv2 (and optionally SAM2) on *pil_image*.
+    """Run OWLv2 (+ optional SAM2) with explicit device/dtype.
 
-    Returns detections enriched with 'mask' and 'segmentation' when
-    mode == "Segmentation".
+    In ZeroGPU mode this is called inside the @spaces.GPU context so CUDA is
+    available; locally it uses whatever settings.device resolved to.
     """
-    processor, model = _get_owlv2()
-    detections = _owlv2_infer(
-        pil_image, processor, model, prompts, threshold,
-        settings.device, settings.torch_dtype,
-    )
+    processor, owlv2 = _get_owlv2()
+    owlv2.to(device)
+    try:
+        detections = _owlv2_infer(
+            pil_image, processor, owlv2, prompts, threshold, device, dtype,
+        )
+    finally:
+        if _ZERO_GPU:
+            owlv2.to("cpu")  # release VRAM back to ZeroGPU pool
 
     if mode == "Segmentation" and detections:
         sam2_processor, sam2_model = _get_sam2()
-        detections = segment_with_boxes(
-            pil_image, detections, sam2_processor, sam2_model, settings.device
-        )
+        sam2_model.to(device)
+        try:
+            detections = segment_with_boxes(
+                pil_image, detections, sam2_processor, sam2_model, device
+            )
+        finally:
+            if _ZERO_GPU:
+                sam2_model.to("cpu")
 
     return detections
 
 
+if _ZERO_GPU:
+    @_spaces.GPU
+    def _run_detection(
+        pil_image: Image.Image,
+        prompts: list[str],
+        threshold: float,
+        mode: str,
+    ) -> list[dict]:
+        """ZeroGPU entry-point: GPU is allocated for the duration of this call."""
+        import torch
+        return _infer_on_device(
+            pil_image, prompts, threshold, mode,
+            device="cuda", dtype=torch.float16,
+        )
+else:
+    def _run_detection(
+        pil_image: Image.Image,
+        prompts: list[str],
+        threshold: float,
+        mode: str,
+    ) -> list[dict]:
+        return _infer_on_device(
+            pil_image, prompts, threshold, mode,
+            device=settings.device, dtype=settings.torch_dtype,
+        )
+
+
 def _parse_prompts(text: str) -> list[str]:
     return [p.strip() for p in text.split(",") if p.strip()]
 

requirements.txt

@@ -11,3 +11,4 @@ click>=8.1.7
 tqdm>=4.66.0
 python-dotenv>=1.0.1
 gradio>=6.0.0
+spaces

samples/CREDITS.txt

@@ -1,4 +1,4 @@
-Sample images used in the autolabel demo
+Cool Sample images used in the autolabel demo
 =========================================
 
 kitchen.jpg