Add ZeroGPU support with @spaces.GPU decorators

README.md

@@ -7,6 +7,7 @@ sdk: gradio
 sdk_version: 5.49.1
 app_file: app.py
 python_version: 3.11
+hardware: zero-gpu
 pinned: false
 ---
 

app.py

@@ -16,7 +16,7 @@ from matplotlib import cm
 from glob import glob
 from natsort import natsorted
 from huggingface_hub import HfApi, upload_file
-# import spaces
+import spaces
 
 from inference_seg import load_model as load_seg_model, run as run_seg
 from inference_count import load_model as load_count_model, run as run_count
@@ -293,7 +293,7 @@ def cleanup_tracking_cache(track_vis_cache):
                 pass
 
 
-# @spaces.GPU
+@spaces.GPU
 def segment_with_choice(use_box_choice, annot_value, overlay_alpha):
     """Segmentation handler - supports bounding box, returns colorized overlay and original mask path"""
     if annot_value is None or len(annot_value) < 1:
@@ -313,7 +313,8 @@ def segment_with_choice(use_box_choice, annot_value, overlay_alpha):
 
 
     try:
-        mask = run_seg(SEG_MODEL, img_path, box=box_array, device=SEG_DEVICE)
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        mask = run_seg(SEG_MODEL, img_path, box=box_array, device=device)
         print("📏 mask shape:", mask.shape, "dtype:", mask.dtype)
     except Exception as e:
         print(f"❌ Inference failed: {str(e)}")
@@ -353,7 +354,7 @@ def segment_with_choice(use_box_choice, annot_value, overlay_alpha):
     return overlay_img, temp_mask_file.name, seg_vis_cache
 
 
-# @spaces.GPU
+@spaces.GPU
 def count_cells_handler(use_box_choice, annot_value, overlay_alpha):
     """Counting handler - supports bounding box, returns only density map"""
     if annot_value is None or len(annot_value) < 1:
@@ -373,11 +374,12 @@ def count_cells_handler(use_box_choice, annot_value, overlay_alpha):
     try:
         print(f"🔢 Counting - Image: {image_path}")
         
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         result = run_count(
             COUNT_MODEL,
             image_path,
             box=box_array,
-            device=COUNT_DEVICE,
+            device=device,
             visualize=True
         )
         
@@ -722,7 +724,7 @@ def create_tracking_visualization(tif_dir, output_dir, valid_tif_files, overlay_
         except:
             return valid_tif_files[0]
 
-# @spaces.GPU
+@spaces.GPU
 def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj, overlay_alpha, prev_track_vis_cache):
     """
     Tracking handler - processes a ZIP of TIF frames, supports bounding box, returns visualization and results ZIP
@@ -808,11 +810,12 @@ def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj, overlay
         print(f"💾 CTC-format results will be saved to: {output_temp_dir}")
 
         # Run tracking with optional bounding box
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         result = run_track(
             TRACK_MODEL,
             video_dir=tif_dir,
             box=box_array,  # Pass bounding box if specified
-            device=TRACK_DEVICE,
+            device=device,
             output_dir=output_temp_dir
         )
         

requirements.txt

@@ -32,6 +32,7 @@ numpy==1.24.4
 # Gradio
 gradio
 gradio-bbox-annotator
+spaces
 
 # Utilities
 natsort