put predictor in session state

app.py

@@ -72,22 +72,24 @@ OBJ_ID = 0
 
 
 @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
+def get_predictor(session_state):
+    if "predictor" not in session_state:
+        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 demo
+        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
 
-predictor = build_predictor()
+        session_state["predictor"] = predictor
+    return session_state["predictor"] = predictor
 
 
 def get_video_fps(video_path):
@@ -106,6 +108,7 @@ def get_video_fps(video_path):
 
 @spaces.GPU
 def reset(session_state):
+    predictor = get_predictor(session_state)
     predictor.to("cuda")
     session_state["input_points"] = []
     session_state["input_labels"] = []
@@ -126,6 +129,7 @@ def reset(session_state):
 
 @spaces.GPU
 def clear_points(session_state):
+    predictor = get_predictor(session_state)
     predictor.to("cuda")
     session_state["input_points"] = []
     session_state["input_labels"] = []
@@ -141,6 +145,7 @@ def clear_points(session_state):
 
 @spaces.GPU
 def preprocess_video_in(video_path, session_state):
+    predictor = get_predictor(session_state)
     predictor.to("cuda")
     if video_path is None:
         return (
@@ -205,6 +210,7 @@ def segment_with_points(
     session_state,
     evt: gr.SelectData,
 ):
+    predictor = get_predictor(session_state)
     predictor.to("cuda")
     session_state["input_points"].append(evt.index)
     print(f"TRACKING INPUT POINT: {session_state['input_points']}")
@@ -279,6 +285,7 @@ def propagate_to_all(
     video_in,
     session_state,
 ):
+    predictor = get_predictor(session_state)
     predictor.to("cuda")
     if (
         len(session_state["input_points"]) == 0