Update app.py

app.py

@@ -320,6 +320,7 @@ def clear_ref_clicks(ref_state):
 
 
 @spaces.GPU(duration=40)
+@torch.no_grad()
 def track_video(n_frames, video_state):
     input_points = video_state["input_points"]
     input_labels = video_state["input_labels"]
@@ -345,48 +346,56 @@ def track_video(n_frames, video_state):
 
     sam2_checkpoint = "./sam2/SAM2-Video-Predictor/checkpoints/sam2_hiera_large.pt"
     config = "sam2_hiera_l.yaml"
-    video_predictor_local = build_sam2_video_predictor(
-        config, sam2_checkpoint, device="cuda"
-    )
 
-    inference_state = video_predictor_local.init_state(
-        images=images_np / 255, device="cuda"
-    )
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        video_predictor_local = build_sam2_video_predictor(
+            config, sam2_checkpoint, device="cuda"
+        )
 
-    if len(torch.from_numpy(video_state["masks"][0]).shape) == 3:
-        mask0 = torch.from_numpy(video_state["masks"][0])[:, :, 0]
-    else:
-        mask0 = torch.from_numpy(video_state["masks"][0])
+        inference_state = video_predictor_local.init_state(
+            images=images_np / 255, device="cuda"
+        )
 
-    video_predictor_local.add_new_mask(
-        inference_state=inference_state,
-        frame_idx=0,
-        obj_id=obj_id,
-        mask=mask0,
-    )
+        if len(torch.from_numpy(video_state["masks"][0]).shape) == 3:
+            mask0 = torch.from_numpy(video_state["masks"][0])[:, :, 0]
+        else:
+            mask0 = torch.from_numpy(video_state["masks"][0])
 
-    output_frames = []
-    mask_frames = []
-    color = (
-        np.array(COLOR_PALETTE[int(time.time()) % len(COLOR_PALETTE)], dtype=np.float32)
-        / 255.0
-    )
-    color = color[None, None, :]
-    for out_frame_idx, out_obj_ids, out_mask_logits in video_predictor_local.propagate_in_video(
-        inference_state
-    ):
-        frame = images_np[out_frame_idx].astype(np.float32) / 255.0
-        mask = np.zeros((H, W, 3), dtype=np.float32)
-        for i, logit in enumerate(out_mask_logits):
-            out_mask = logit.cpu().squeeze().detach().numpy()
-            out_mask = (out_mask[:, :, None] > 0).astype(np.float32)
-            mask += out_mask
-        mask = np.clip(mask, 0, 1)
-        mask = cv2.resize(mask, (W_, H_))
-        mask_frames.append(mask)
-        painted = (1 - mask * 0.5) * frame + mask * 0.5 * color
-        painted = np.uint8(np.clip(painted * 255, 0, 255))
-        output_frames.append(painted)
+        video_predictor_local.add_new_mask(
+            inference_state=inference_state,
+            frame_idx=0,
+            obj_id=obj_id,
+            mask=mask0,
+        )
+
+        output_frames = []
+        mask_frames = []
+        color = (
+            np.array(
+                COLOR_PALETTE[int(time.time()) % len(COLOR_PALETTE)],
+                dtype=np.float32,
+            )
+            / 255.0
+        )
+        color = color[None, None, :]
+
+        for out_frame_idx, out_obj_ids, out_mask_logits in video_predictor_local.propagate_in_video(
+            inference_state
+        ):
+            frame = images_np[out_frame_idx].astype(np.float32) / 255.0
+
+            mask = np.zeros((H, W, 3), dtype=np.float32)
+            for i, logit in enumerate(out_mask_logits):
+                out_mask = logit.cpu().squeeze().detach().numpy()
+                out_mask = (out_mask[:, :, None] > 0).astype(np.float32)
+                mask += out_mask
+
+            mask = np.clip(mask, 0, 1)
+            mask = cv2.resize(mask, (W_, H_))
+            mask_frames.append(mask)
+            painted = (1 - mask * 0.5) * frame + mask * 0.5 * color
+            painted = np.uint8(np.clip(painted * 255, 0, 255))
+            output_frames.append(painted)
 
     video_file = f"/tmp/{time.time()}-{random.random()}-tracked_output.mp4"
     clip = ImageSequenceClip(output_frames, fps=15)