make gr.State individual

app.py

@@ -90,16 +90,22 @@ def get_video_fps(video_path):
     return fps
 
 
-def reset():
+def reset(
+    session_first_frame,
+    session_all_frames,
+    session_input_points,
+    session_input_labels,
+    request: gr.Request,
+):
+    session_id = request.session_id
     predictor.to("cpu")
-    session_state["input_points"] = []
-    session_state["input_labels"] = []
+    session_input_points = []
+    session_input_labels = []
 
-    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_first_frame = None
+    session_all_frames = None
     global_inference_states[session_id] = None
     return (
         None,
@@ -107,26 +113,38 @@ def reset():
         None,
         None,
         gr.update(value=None, visible=False),
-        session_state,
+        session_first_frame,
+        session_all_frames,
+        session_input_points,
+        session_input_labels,
     )
 
 
-def clear_points(session_state):
+def clear_points(session_input_points, session_input_labels, request: gr.Request,):
+    session_id = request.session_id
     predictor.to("cpu")
-    session_state["input_points"] = []
-    session_state["input_labels"] = []
-    session_id = id(session_state)
+    session_input_points = []
+    session_input_labels = []
     if global_inference_states[session_id]["tracking_has_started"]:
         predictor.reset_state(global_inference_states[session_id])
     return (
-        session_state["first_frame"],
+        session_first_frame,
         None,
         gr.update(value=None, visible=False),
-        session_state,
+        session_input_points,
+        session_input_labels,
     )
 
 
-def preprocess_video_in(video_path, session_state):
+def preprocess_video_in(
+    video_path,
+    session_first_frame,
+    session_all_frames,
+    session_input_points,
+    session_input_labels,
+    request: gr.Request,
+):
+    session_id = request.session_id
     predictor.to("cpu")
     if video_path is None:
         return (
@@ -134,7 +152,10 @@ def preprocess_video_in(video_path, session_state):
             None,  # points_map
             None,  # output_image
             gr.update(value=None, visible=False),  # output_video
-            session_state,
+            session_first_frame,
+            session_all_frames,
+            session_input_points,
+            session_input_labels,
         )
 
     # Read the first frame
@@ -146,7 +167,10 @@ def preprocess_video_in(video_path, session_state):
             None,  # points_map
             None,  # output_image
             gr.update(value=None, visible=False),  # output_video
-            session_state,
+            session_first_frame,
+            session_all_frames,
+            session_input_points,
+            session_input_labels,
         )
 
     frame_number = 0
@@ -169,46 +193,51 @@ def preprocess_video_in(video_path, session_state):
         frame_number += 1
 
     cap.release()
-    session_state["first_frame"] = copy.deepcopy(first_frame)
-    session_state["all_frames"] = all_frames
+    session_first_frame = copy.deepcopy(first_frame)
+    session_all_frames = all_frames
 
-    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"] = []
+    session_input_points = []
+    session_input_labels = []
 
     return [
         gr.update(open=False),  # video_in_drawer
         first_frame,  # points_map
         None,  # output_image
         gr.update(value=None, visible=False),  # output_video
-        session_state,
+        session_first_frame,
+        session_all_frames,
+        session_input_points,
+        session_input_labels,
     ]
 
 
 @spaces.GPU
 def segment_with_points(
     point_type,
-    session_state,
+    session_input_points,
+    session_input_labels,
     evt: gr.SelectData,
+    request: gr.Request,
 ):
+    session_id = request.session_id
     if torch.cuda.get_device_properties(0).major >= 8:
         torch.backends.cuda.matmul.allow_tf32 = True
         torch.backends.cudnn.allow_tf32 = True
     with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
         predictor.to("cuda")
-        session_state["input_points"].append(evt.index)
-        print(f"TRACKING INPUT POINT: {session_state['input_points']}")
+        session_input_points.append(evt.index)
+        print(f"TRACKING INPUT POINT: {session_input_points}")
 
         if point_type == "include":
-            session_state["input_labels"].append(1)
+            session_input_labels.append(1)
         elif point_type == "exclude":
-            session_state["input_labels"].append(0)
-        print(f"TRACKING INPUT LABEL: {session_state['input_labels']}")
+            session_input_labels.append(0)
+        print(f"TRACKING INPUT LABEL: {session_input_labels}")
 
         # Open the image and get its dimensions
-        transparent_background = Image.fromarray(session_state["first_frame"]).convert(
+        transparent_background Image.fromarray(session_first_frame).convert(
             "RGBA"
         )
         w, h = transparent_background.size
@@ -220,8 +249,8 @@ def segment_with_points(
         # Create a transparent layer to draw on
         transparent_layer = np.zeros((h, w, 4), dtype=np.uint8)
 
-        for index, track in enumerate(session_state["input_points"]):
-            if session_state["input_labels"][index] == 1:
+        for index, track in enumerate(session_input_points):
+            if session_input_labels[index] == 1:
                 cv2.circle(transparent_layer, track, radius, (0, 255, 0, 255), -1)
             else:
                 cv2.circle(transparent_layer, track, radius, (255, 0, 0, 255), -1)
@@ -233,10 +262,9 @@ def segment_with_points(
         )
 
         # Let's add a positive click at (x, y) = (210, 350) to get started
-        points = np.array(session_state["input_points"], dtype=np.float32)
+        points = np.array(session_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)
+        labels = np.array(session_input_labels, dtype=np.int32)
         _, _, out_mask_logits = predictor.add_new_points(
             inference_state=global_inference_states[session_id],
             frame_idx=0,
@@ -249,7 +277,7 @@ def segment_with_points(
         first_frame_output = Image.alpha_composite(transparent_background, mask_image)
 
         torch.cuda.empty_cache()
-        return selected_point_map, first_frame_output, session_state
+        return selected_point_map, first_frame_output, session_input_points, session_input_labels
 
 
 def show_mask(mask, obj_id=None, random_color=False, convert_to_image=True):
@@ -270,23 +298,21 @@ def show_mask(mask, obj_id=None, random_color=False, convert_to_image=True):
 @spaces.GPU
 def propagate_to_all(
     video_in,
-    session_state,
+    session_all_frames,
+    request: gr.Request,
 ):
+    session_id = request.session_id
     predictor.to("cuda")
     if torch.cuda.get_device_properties(0).major >= 8:
         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
+            len (session_input_points) == 0
             or video_in is None
             or global_inference_states[session_id] is None
         ):
-            return (
-                None,
-                session_state,
-            )
+            return None
 
         # run propagation throughout the video and collect the results in a dict
         video_segments = (
@@ -307,7 +333,7 @@ def propagate_to_all(
         output_frames = []
         for out_frame_idx in range(0, len(video_segments), vis_frame_stride):
             transparent_background = Image.fromarray(
-                session_state["all_frames"][out_frame_idx]
+                session_all_frames[out_frame_idx]
             ).convert("RGBA")
             out_mask = video_segments[out_frame_idx][OBJ_ID]
             mask_image = show_mask(out_mask)
@@ -331,10 +357,7 @@ def propagate_to_all(
         # Write the result to a file
         clip.write_videofile(final_vid_output_path, codec="libx264")
 
-        return (
-            gr.update(value=final_vid_output_path),
-            session_state,
-        )
+        return gr.update(value=final_vid_output_path)
 
 
 def update_ui():
@@ -342,14 +365,10 @@ def update_ui():
 
 
 with gr.Blocks() as demo:
-    session_state = gr.State(
-        {
-            "first_frame": None,
-            "all_frames": None,
-            "input_points": [],
-            "input_labels": [],
-        }
-    )
+    first_frame = gr.State(None)
+    all_frames = gr.State(None)
+    input_points = gr.State([])
+    input_labels = gr.State([])
 
     with gr.Column():
         # Title
@@ -399,14 +418,20 @@ with gr.Blocks() as demo:
         fn=preprocess_video_in,
         inputs=[
             video_in,
-            session_state,
+            first_frame,
+            all_frames,
+            input_points,
+            input_labels,
         ],
         outputs=[
             video_in_drawer,  # Accordion to hide uploaded video player
             points_map,  # Image component where we add new tracking points
             output_image,
             output_video,
-            session_state,
+            first_frame,
+            all_frames,
+            input_points,
+            input_labels,
         ],
         queue=False,
     )
@@ -415,14 +440,20 @@ with gr.Blocks() as demo:
         fn=preprocess_video_in,
         inputs=[
             video_in,
-            session_state,
+            first_frame,
+            all_frames,
+            input_points,
+            input_labels,
         ],
         outputs=[
             video_in_drawer,  # Accordion to hide uploaded video player
             points_map,  # Image component where we add new tracking points
             output_image,
             output_video,
-            session_state,
+            first_frame,
+            all_frames,
+            input_points,
+            input_labels,
         ],
         queue=False,
     )
@@ -432,12 +463,14 @@ with gr.Blocks() as demo:
         fn=segment_with_points,
         inputs=[
             point_type,  # "include" or "exclude"
-            session_state,
+            input_points,
+            input_labels,
         ],
         outputs=[
             points_map,  # updated image with points
             output_image,
-            session_state,
+            input_points,
+            input_labels,
         ],
         queue=False,
     )
@@ -445,26 +478,38 @@ with gr.Blocks() as demo:
     # Clear every points clicked and added to the map
     clear_points_btn.click(
         fn=clear_points,
-        inputs=session_state,
+        inputs=[
+            input_points,
+            input_labels,
+        ],
         outputs=[
             points_map,
             output_image,
             output_video,
-            session_state,
+            input_points,
+            input_labels,
         ],
         queue=False,
     )
 
     reset_btn.click(
         fn=reset,
-        inputs=session_state,
+        inputs=[
+            first_frame,
+            all_frames,
+            input_points,
+            input_labels,
+        ],
         outputs=[
             video_in,
             video_in_drawer,
             points_map,
             output_image,
             output_video,
-            session_state,
+            first_frame,
+            all_frames,
+            input_points,
+            input_labels,
         ],
         queue=False,
     )
@@ -478,11 +523,10 @@ with gr.Blocks() as demo:
         fn=propagate_to_all,
         inputs=[
             video_in,
-            session_state,
+            all_frames,
         ],
         outputs=[
             output_video,
-            session_state,
         ],
         concurrency_limit=10,
         queue=False,