rollback

.gitignore

@@ -0,0 +1,2 @@
+*.egg-info/
+__pycache__/

app.py

@@ -71,25 +71,17 @@ examples = [
 OBJ_ID = 0
 
 
-@spaces.GPU
-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
+sam2_checkpoint = "checkpoints/edgetam.pt"
+model_cfg = "edgetam.yaml"
+predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device="cuda")
+print("predictor loaded")
 
-        session_state["predictor"] = predictor
-    return session_state["predictor"]
+# 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
 
 
 def get_video_fps(video_path):
@@ -106,10 +98,7 @@ def get_video_fps(video_path):
     return fps
 
 
-@spaces.GPU
 def reset(session_state):
-    predictor = get_predictor(session_state)
-    predictor.to("cuda")
     session_state["input_points"] = []
     session_state["input_labels"] = []
     if session_state["inference_state"] is not None:
@@ -127,10 +116,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"] = []
     if session_state["inference_state"]["tracking_has_started"]:
@@ -145,8 +131,6 @@ 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 (
             gr.update(open=True),  # video_in_drawer
@@ -210,8 +194,6 @@ 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']}")
 
@@ -285,8 +267,6 @@ def propagate_to_all(
     video_in,
     session_state,
 ):
-    predictor = get_predictor(session_state)
-    predictor.to("cuda")
     if (
         len(session_state["input_points"]) == 0
         or video_in is None