zerogpu hooks

app.py

@@ -4,6 +4,7 @@ import numpy as np
 import pandas as pd
 from PIL import Image
 import torch
+import spaces
 
 from mussel.utils import get_features, segment_tissue
 from mussel.models import ModelType
@@ -11,21 +12,11 @@ from aeon import run_aeon
 from paladin import run_paladin
 
 NUM_WORKERS = int(os.getenv("NUM_WORKERS", 2))
-USE_GPU = torch.cuda.is_available()
 
-def analyze_slide(slide_path, site_type):
-    _, _, coords, attrs = segment_tissue(
-        slide_path=slide_path,
-        patch_size=224,
-        segment_threshold=15,
-        median_blur_ksize=11,
-        morphology_ex_kernel=2,
-        tissue_area_threshold=2,
-        hole_area_threshold=1,
-        max_num_holes=2
-    )
-
-    features, _ = get_features(coords, slide_path, attrs, model_type=ModelType.OPTIMUS, use_gpu=USE_GPU, batch_size=64, num_workers=NUM_WORKERS)
+@spaces.GPU
+def get_features_and_infer(coords, slide_path, attrs, site_type):
+    use_gpu = torch.cuda.is_available()
+    features, _ = get_features(coords, slide_path, attrs, model_type=ModelType.OPTIMUS, use_gpu=use_gpu, batch_size=64, num_workers=NUM_WORKERS)
 
     # Step 3: Run Aeon to predict histology
     aeon_results, _ = run_aeon(
@@ -33,8 +24,8 @@ def analyze_slide(slide_path, site_type):
         model_path="data/aeon_model.pkl",
         metastatic=(site_type == "Metastatic"),
         batch_size=8,
-        num_workers=NM_WORKERS,
-        use_cpu=not USE_GPU,
+        num_workers=NUM_WORKERS,
+        use_cpu=not use_gpu,
     )
 
     # Step 4: Run Paladin to predict biomarkers
@@ -45,11 +36,27 @@ def analyze_slide(slide_path, site_type):
         metastatic=(site_type == "Metastatic"),
         batch_size=8,
         num_workers=NUM_WORKERS,
-        use_cpu=not USE_GPU,
+        use_cpu=not use_gpu,
     )
 
     return aeon_results, paladin_results
 
+
+def analyze_slide(slide_path, site_type):
+    _, _, coords, attrs = segment_tissue(
+        slide_path=slide_path,
+        patch_size=224,
+        segment_threshold=15,
+        median_blur_ksize=11,
+        morphology_ex_kernel=2,
+        tissue_area_threshold=2,
+        hole_area_threshold=1,
+        max_num_holes=2
+    )
+    aeon_results, paladin_results = get_features_and_infer(coords, slide_path, attrs, site_type)
+    return aeon_results, paladin_results
+
+
 if __name__ == "__main__":
     input_slide = gr.File(
         label="Upload H&E Whole Slide Image",