move inference_states out of gr.State

app.py

@@ -73,6 +73,7 @@ OBJ_ID = 0
 sam2_checkpoint = "checkpoints/edgetam.pt"
 model_cfg = "edgetam.yaml"
 predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device="cpu")
+global_inference_states = {}
 
 
 def get_video_fps(video_path):
@@ -89,15 +90,17 @@ def get_video_fps(video_path):
     return fps
 
 
-def reset(session_state):
+def reset():
     predictor.to("cpu")
     session_state["input_points"] = []
     session_state["input_labels"] = []
-    if session_state["inference_state"] is not None:
-        predictor.reset_state(session_state["inference_state"])
+
+    session_id = id(session_state)
+    if global_inference_states[session_id] is not None:
+        predictor.reset_state(global_inference_states[session_id])
     session_state["first_frame"] = None
     session_state["all_frames"] = None
-    session_state["inference_state"] = None
+    global_inference_states[session_id] = None
     return (
         None,
         gr.update(open=True),
@@ -112,8 +115,9 @@ def clear_points(session_state):
     predictor.to("cpu")
     session_state["input_points"] = []
     session_state["input_labels"] = []
-    if session_state["inference_state"]["tracking_has_started"]:
-        predictor.reset_state(session_state["inference_state"])
+    session_id = id(session_state)
+    if global_inference_states[session_id]["tracking_has_started"]:
+        predictor.reset_state(global_inference_states[session_id])
     return (
         session_state["first_frame"],
         None,
@@ -168,7 +172,9 @@ def preprocess_video_in(video_path, session_state):
     session_state["first_frame"] = copy.deepcopy(first_frame)
     session_state["all_frames"] = all_frames
 
-    session_state["inference_state"] = predictor.init_state(video_path=video_path)
+    session_id = id(session_state)
+    global_inference_states[session_id] = predictor.init_state(video_path=video_path)
+
     session_state["input_points"] = []
     session_state["input_labels"] = []
 
@@ -230,8 +236,9 @@ def segment_with_points(
         points = np.array(session_state["input_points"], dtype=np.float32)
         # for labels, `1` means positive click and `0` means negative click
         labels = np.array(session_state["input_labels"], np.int32)
+        session_id = id(session_state)
         _, _, out_mask_logits = predictor.add_new_points(
-            inference_state=session_state["inference_state"],
+            inference_state=global_inference_states[session_id],
             frame_idx=0,
             obj_id=OBJ_ID,
             points=points,
@@ -270,10 +277,11 @@ def propagate_to_all(
         torch.backends.cuda.matmul.allow_tf32 = True
         torch.backends.cudnn.allow_tf32 = True
     with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        session_id = id(session_state)
         if (
             len(session_state["input_points"]) == 0
             or video_in is None
-            or session_state["inference_state"] is None
+            or global_inference_states[session_id] is None
         ):
             return (
                 None,
@@ -286,7 +294,7 @@ def propagate_to_all(
         )  # video_segments contains the per-frame segmentation results
         print("starting propagate_in_video")
         for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(
-            session_state["inference_state"]
+            global_inference_states[session_id]
         ):
             video_segments[out_frame_idx] = {
                 out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
@@ -340,7 +348,6 @@ with gr.Blocks() as demo:
             "all_frames": None,
             "input_points": [],
             "input_labels": [],
-            "inference_state": None,
         }
     )