remove CUDA init in main process

app.py

@@ -69,21 +69,27 @@ examples = [
 ]
 
 OBJ_ID = 0
-sam2_checkpoint = "checkpoints/edgetam.pt"
-model_cfg = "edgetam.yaml"
-predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device="cuda")
-print("PREDICTOR LOADED")
-predictor.to("cuda")
-
-# use bfloat16 for the entire notebook
-torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
-if torch.cuda.get_device_properties(0).major >= 8:
-    # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
-    torch.backends.cuda.matmul.allow_tf32 = True
-    torch.backends.cudnn.allow_tf32 = True
 
 
 @spaces.GPU
+def build_predictor():
+    sam2_checkpoint = "checkpoints/edgetam.pt"
+    model_cfg = "edgetam.yaml"
+    predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device="cuda")
+    print("predictor loaded")
+
+    # use bfloat16 for the entire notebook
+    torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
+    if torch.cuda.get_device_properties(0).major >= 8:
+        # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+    return predictor
+
+
+predictor = build_predictor()
+
+
 def get_video_fps(video_path):
     # Open the video file
     cap = cv2.VideoCapture(video_path)
@@ -100,6 +106,7 @@ def get_video_fps(video_path):
 
 @spaces.GPU
 def reset(session_state):
+    predictor.to("cuda")
     session_state["input_points"] = []
     session_state["input_labels"] = []
     if session_state["inference_state"] is not None:
@@ -119,6 +126,7 @@ def reset(session_state):
 
 @spaces.GPU
 def clear_points(session_state):
+    predictor.to("cuda")
     session_state["input_points"] = []
     session_state["input_labels"] = []
     if session_state["inference_state"]["tracking_has_started"]:
@@ -133,6 +141,7 @@ def clear_points(session_state):
 
 @spaces.GPU
 def preprocess_video_in(video_path, session_state):
+    predictor.to("cuda")
     if video_path is None:
         return (
             gr.update(open=True),  # video_in_drawer
@@ -196,6 +205,7 @@ def segment_with_points(
     session_state,
     evt: gr.SelectData,
 ):
+    predictor.to("cuda")
     session_state["input_points"].append(evt.index)
     print(f"TRACKING INPUT POINT: {session_state['input_points']}")
 
@@ -249,7 +259,6 @@ def segment_with_points(
     return selected_point_map, first_frame_output, session_state
 
 
-@spaces.GPU
 def show_mask(mask, obj_id=None, random_color=False, convert_to_image=True):
     if random_color:
         color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
@@ -270,6 +279,7 @@ def propagate_to_all(
     video_in,
     session_state,
 ):
+    predictor.to("cuda")
     if (
         len(session_state["input_points"]) == 0
         or video_in is None
@@ -325,7 +335,6 @@ def propagate_to_all(
     )
 
 
-@spaces.GPU
 def update_ui():
     return gr.update(visible=True)
 
@@ -478,5 +487,6 @@ with gr.Blocks() as demo:
         queue=False,
     )
 
+
 demo.queue()
 demo.launch()