quoccuongle/Grounded-SAM-2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

patch for the case of offload_state_to_cpu=True in the new SAM2VideoPredictor (#490)

Browse Source 
This PR adds a pathc for the case of `offload_state_to_cpu=True` where `pred_masks` might have been offload to CPU device (close https://github.com/facebookresearch/sam2/issues/489)
This commit is contained in:
Ronghang Hu
2024-12-12 15:12:13 -08:00
committed by GitHub
parent 393ae336a7
commit 722d1d1511
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
sam2/sam2_video_predictor.py
View File

@@ -591,7 +591,8 @@ class SAM2VideoPredictor(SAM2Base):
if frame_idx in obj_output_dict["cond_frame_outputs"]: if frame_idx in obj_output_dict["cond_frame_outputs"]:
storage_key = "cond_frame_outputs" storage_key = "cond_frame_outputs"
current_out = obj_output_dict[storage_key][frame_idx] current_out = obj_output_dict[storage_key][frame_idx]
pred_masks = current_out["pred_masks"] device = inference_state["device"]
pred_masks = current_out["pred_masks"].to(device, non_blocking=True)
if self.clear_non_cond_mem_around_input: if self.clear_non_cond_mem_around_input:
# clear non-conditioning memory of the surrounding frames # clear non-conditioning memory of the surrounding frames
self._clear_obj_non_cond_mem_around_input( self._clear_obj_non_cond_mem_around_input(