rollback sam2_base and sam2_video_predictor

sam2/modeling/sam2_base.py

@@ -617,7 +617,7 @@ class SAM2Base(torch.nn.Module):
                             if self.use_signed_tpos_enc_to_obj_ptrs
                             else abs(frame_idx - t)
                         ),
-                        out["obj_ptr"].to(device),
+                        out["obj_ptr"],
                     )
                     for t, out in ptr_cond_outputs.items()
                 ]
@@ -630,7 +630,7 @@ class SAM2Base(torch.nn.Module):
                         t, unselected_cond_outputs.get(t, None)
                     )
                     if out is not None:
-                        pos_and_ptrs.append((t_diff, out["obj_ptr"].to(device)))
+                        pos_and_ptrs.append((t_diff, out["obj_ptr"]))
                 # If we have at least one object pointer, add them to the across attention
                 if len(pos_and_ptrs) > 0:
                     pos_list, ptrs_list = zip(*pos_and_ptrs)

sam2/sam2_video_predictor.py

@@ -470,7 +470,7 @@ class SAM2VideoPredictor(SAM2Base):
                 size=(batch_size, self.hidden_dim),
                 fill_value=NO_OBJ_SCORE,
                 dtype=torch.float32,
-                device=inference_state["storage_device"],
+                device=inference_state["device"],
             ),
             "object_score_logits": torch.full(
                 size=(batch_size, 1),
@@ -478,7 +478,7 @@ class SAM2VideoPredictor(SAM2Base):
                 # present as sigmoid(10)=1, same as in `predict_masks` of `MaskDecoder`
                 fill_value=10.0,
                 dtype=torch.float32,
-                device=inference_state["storage_device"],
+                device=inference_state["device"],
             ),
         }
         empty_mask_ptr = None
@@ -545,9 +545,7 @@ class SAM2VideoPredictor(SAM2Base):
                 frame_idx=frame_idx,
                 batch_size=batch_size,
                 high_res_masks=high_res_masks,
-                object_score_logits=consolidated_out["object_score_logits"].to(
-                    device, non_blocking=True
-                ),
+                object_score_logits=consolidated_out["object_score_logits"],
                 is_mask_from_pts=True,  # these frames are what the user interacted with
             )
             consolidated_out["maskmem_features"] = maskmem_features
@@ -881,10 +879,9 @@ class SAM2VideoPredictor(SAM2Base):
     def _get_image_feature(self, inference_state, frame_idx, batch_size):
         """Compute the image features on a given frame."""
         # Look up in the cache first
-        # image, backbone_out = inference_state["cached_features"].get(
-        #     frame_idx, (None, None)
-        # )
-        image, backbone_out = None, None
+        image, backbone_out = inference_state["cached_features"].get(
+            frame_idx, (None, None)
+        )
         if backbone_out is None:
             # Cache miss -- we will run inference on a single image
             device = inference_state["device"]
@@ -892,7 +889,7 @@ class SAM2VideoPredictor(SAM2Base):
             backbone_out = self.forward_image(image)
             # Cache the most recent frame's feature (for repeated interactions with
             # a frame; we can use an LRU cache for more frames in the future).
-            # inference_state["cached_features"] = {frame_idx: (image, backbone_out)}
+            inference_state["cached_features"] = {frame_idx: (image, backbone_out)}
 
         # expand the features to have the same dimension as the number of objects
         expanded_image = image.expand(batch_size, -1, -1, -1)
@@ -967,11 +964,9 @@ class SAM2VideoPredictor(SAM2Base):
         pred_masks = pred_masks_gpu.to(storage_device, non_blocking=True)
         # "maskmem_pos_enc" is the same across frames, so we only need to store one copy of it
         maskmem_pos_enc = self._get_maskmem_pos_enc(inference_state, current_out)
-        # object pointer is a small tensor, so we always keep it on GPU memory for fast access (modified for ZeroGPU)
-        obj_ptr = current_out["obj_ptr"].to(storage_device, non_blocking=True)
-        object_score_logits = current_out["object_score_logits"].to(
-            storage_device, non_blocking=True
-        )
+        # object pointer is a small tensor, so we always keep it on GPU memory for fast access
+        obj_ptr = current_out["obj_ptr"]
+        object_score_logits = current_out["object_score_logits"]
         # make a compact version of this frame's output to reduce the state size
         compact_current_out = {
             "maskmem_features": maskmem_features,
@@ -1023,7 +1018,6 @@ class SAM2VideoPredictor(SAM2Base):
         `maskmem_pos_enc` is the same across frames and objects, so we cache it as
         a constant in the inference session to reduce session storage size.
         """
-        storage_device = inference_state["storage_device"]
         model_constants = inference_state["constants"]
         # "out_maskmem_pos_enc" should be either a list of tensors or None
         out_maskmem_pos_enc = current_out["maskmem_pos_enc"]
@@ -1032,9 +1026,7 @@ class SAM2VideoPredictor(SAM2Base):
                 assert isinstance(out_maskmem_pos_enc, list)
                 # only take the slice for one object, since it's same across objects
                 maskmem_pos_enc = [x[0:1].clone() for x in out_maskmem_pos_enc]
-                model_constants["maskmem_pos_enc"] = maskmem_pos_enc.to(
-                    storage_device, non_blocking=True
-                )
+                model_constants["maskmem_pos_enc"] = maskmem_pos_enc
             else:
                 maskmem_pos_enc = model_constants["maskmem_pos_enc"]
             # expand the cached maskmem_pos_enc to the actual batch size

trimm_examples.py

@@ -0,0 +1,30 @@
+import os
+
+from moviepy.editor import VideoFileClip
+
+# Define the folder and duration
+input_folder = "examples"
+output_folder = "examples/trimmed"
+trim_duration = 3  # seconds
+
+# Create output folder if it doesn't exist
+os.makedirs(output_folder, exist_ok=True)
+
+# Process each .mp4 file
+for filename in os.listdir(input_folder):
+    if filename.lower().endswith(".mp4"):
+        input_path = os.path.join(input_folder, filename)
+        output_path = os.path.join(output_folder, filename)
+
+        print(f"Trimming: {input_path} -> {output_path}")
+        try:
+            clip = VideoFileClip(input_path).subclip(0, trim_duration)
+            clip.write_videofile(
+                output_path,
+                codec="libx264",
+                audio_codec="aac",
+                verbose=False,
+                logger=None,
+            )
+        except Exception as e:
+            print(f"Failed to process {filename}: {e}")