Reduce ZeroGPU occupancy by moving CPU output work outside GPU tasks

app.py

@@ -97,6 +97,7 @@ UPSCALE = 1.0
 
 _model = count.build_pet(torch.device("cpu"))
 _model.load_state_dict(count.load_state_dict(count.DEFAULT_WEIGHTS), strict=True)
+_model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
 
 _logo = Image.open("assets/logo_white.png").convert("RGBA")
 
@@ -126,43 +127,40 @@ def add_watermark(image):
     return image.convert("RGB")
 
 
+@spaces.GPU(duration=7)
+@torch.inference_mode()
+def predict_crowd_points(image_bgr, threshold):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return count.predict_points_tiled(
+        _model, image_bgr, device, UPSCALE, TILE, OVERLAP,
+        float(threshold), time_budget=6.5,
+    )
+
+
 def _count_crowd(image_path, threshold, save_images):
     if image_path is None:
         return None, "Upload an image first."
     image_bgr = cv2.imread(image_path)
     if image_bgr is None:
         return None, "Could not read that image."
-
     try:
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model = _model.to(device)
-        with torch.inference_mode():
-            points = count.predict_points_tiled(
-                model,
-                image_bgr,
-                device,
-                UPSCALE,
-                TILE,
-                OVERLAP,
-                float(threshold),
-                time_budget=6.5,
-            )
+        points = predict_crowd_points(image_bgr, threshold)
         total = int(points.shape[0])
         overlay = add_watermark(count.render_overlay(image_bgr, points, total, logo=_logo))
         if save_images:
             save_images_async(image_path, overlay, total)
         return overlay, f"{total:,} people"
+    except gr.Error:
+        raise
     except Exception as exc:
         print(f"count_crowd failed: {exc}", flush=True)
         return None, "Crowd counter failed. Please try a smaller image or try again in a minute."
 
 
-@spaces.GPU(duration=7)
 def count_crowd(image_path, threshold):
     return _count_crowd(image_path, threshold, save_images=True)
 
 
-@spaces.GPU(duration=7)
 def recount_crowd(image_path, threshold):
     return _count_crowd(image_path, threshold, save_images=False)