Update app.py

app.py

@@ -77,6 +77,9 @@ def clear_points(image):
         image,   # points_map
     ]
 
+@spaces.GPU
+@torch.inference_mode()
+@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def preprocess_video_in(video_path):
     if video_path is None:
         return None, gr.State([]), gr.State([]), None, None, None, None, None, None, gr.update(open=True)
@@ -149,6 +152,9 @@ def preprocess_video_in(video_path):
     ]
 
 
+@spaces.GPU
+@torch.inference_mode()
+@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def get_point(point_type, tracking_points, trackings_input_label, input_first_frame_image, evt: gr.SelectData):
     if input_first_frame_image is None:
         return gr.State([]), gr.State([]), None
@@ -193,7 +199,8 @@ 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 show_mask(mask, ax, obj_id=None, random_color=False):
     if random_color:
         color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
@@ -206,19 +213,20 @@ def show_mask(mask, ax, obj_id=None, random_color=False):
     ax.axis('off')
     ax.imshow(mask_image)
 
-
+@spaces.GPU
 def show_points(coords, labels, ax, marker_size=200):
     pos_points = coords[labels==1]
     neg_points = coords[labels==0]
     ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
     ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
 
+@spaces.GPU
 def show_box(box, ax):
     x0, y0 = box[0], box[1]
     w, h = box[2] - box[0], box[3] - box[1]
     ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))    
 
-
+@spaces.GPU
 def load_model(checkpoint):
     # Load model accordingly to user's choice
     if checkpoint == "efficienttam_s":
@@ -431,12 +439,14 @@ def propagate_to_all(tracking_points, video_in, checkpoint, stored_inference_sta
         
         return gr.update(value=None), gr.update(value=final_vid_output_path), working_frame, available_frames_to_check, gr.update(visible=True)
 
+@spaces.GPU
 def update_ui(vis_frame_type):
     if vis_frame_type == "coarse":
         return gr.update(visible=True), gr.update(visible=False)
     elif vis_frame_type == "fine":
         return gr.update(visible=False), gr.update(visible=True)
 
+@spaces.GPU
 def switch_working_frame(working_frame, scanned_frames, video_frames_dir):
     new_working_frame = None
     if working_frame == None:
@@ -452,6 +462,7 @@ def switch_working_frame(working_frame, scanned_frames, video_frames_dir):
             new_working_frame = os.path.join(video_frames_dir, scanned_frames[ann_frame_idx])
     return gr.State([]), gr.State([]), new_working_frame, new_working_frame
 
+@spaces.GPU
 def reset_propagation(first_frame_path, predictor, stored_inference_state):
     predictor.reset_state(stored_inference_state)
     # print(f"RESET State: {stored_inference_state} ")