Video Segmentation

class Sam3VideoPredictor:
    # Global dictionary holding all inference states
    _ALL_INFERENCE_STATES = {}
    
    def start_session(self, resource_path, session_id=None):
        """Start a new inference session on a video."""
        # Initialize inference state from model
        inference_state = self.model.init_state(
            resource_path=resource_path,
            async_loading_frames=self.async_loading_frames,
            video_loader_type=self.video_loader_type,
        )
        
        if not session_id:
            session_id = str(uuid.uuid4())
        
        self._ALL_INFERENCE_STATES[session_id] = {
            "state": inference_state,
            "session_id": session_id,
            "start_time": time.time(),
        }
        
        return {"session_id": session_id}

def track_step(
    self,
    frame_idx,
    is_init_cond_frame,  # True if this is a prompted frame
    current_vision_feats,
    point_inputs,
    mask_inputs,
    output_dict,
    num_frames,
    run_mem_encoder=True,
):
    # Get high-res features for SAM decoder
    if len(current_vision_feats) > 1:
        high_res_features = [
            x.permute(1, 2, 0).view(x.size(1), x.size(2), *s)
            for x, s in zip(current_vision_feats[:-1], feat_sizes[:-1])
        ]
    
    if mask_inputs is not None:
        # Use mask directly as output
        sam_outputs = self._use_mask_as_output(
            pix_feat, high_res_features, mask_inputs
        )
    else:
        # Condition on memory from previous frames
        pix_feat_with_mem = self._prepare_memory_conditioned_features(
            frame_idx=frame_idx,
            is_init_cond_frame=is_init_cond_frame,
            current_vision_feats=current_vision_feats[-1:],
            output_dict=output_dict,
            num_frames=num_frames,
        )
        
        # Run SAM-style mask decoder
        sam_outputs = self._forward_sam_heads(
            backbone_features=pix_feat_with_mem,
            point_inputs=point_inputs,
            mask_inputs=mask_inputs,
            high_res_features=high_res_features,
            multimask_output=self._use_multimask(is_init_cond_frame, point_inputs),
        )
    
    # Extract masks and object pointers
    low_res_masks, high_res_masks, obj_ptr = sam_outputs[3:6]
    
    # Store in output
    current_out["pred_masks"] = low_res_masks
    current_out["pred_masks_high_res"] = high_res_masks
    current_out["obj_ptr"] = obj_ptr  # Used for memory
    
    return current_out

def _encode_new_memory(
    self,
    image,
    current_vision_feats,
    feat_sizes,
    pred_masks_high_res,
    object_score_logits,
    is_mask_from_pts,
):
    """Encode current frame prediction into memory."""
    B = current_vision_feats[-1].size(1)
    C = self.hidden_dim
    H, W = feat_sizes[-1]
    
    # Get pixel features
    pix_feat = current_vision_feats[-1].permute(1, 2, 0).view(B, C, H, W)
    
    # Apply sigmoid to mask logits
    mask_for_mem = torch.sigmoid(pred_masks_high_res)
    mask_for_mem = mask_for_mem * self.sigmoid_scale_for_mem_enc
    mask_for_mem = mask_for_mem + self.sigmoid_bias_for_mem_enc
    
    # Encode mask + image features
    maskmem_out = self.maskmem_backbone(image, pix_feat, mask_for_mem)
    
    maskmem_features = maskmem_out["vision_features"]
    maskmem_pos_enc = maskmem_out["vision_pos_enc"]
    
    # Add no-object embedding for occluded frames
    is_obj_appearing = (object_score_logits > 0).float()
    maskmem_features += (
        1 - is_obj_appearing[..., None, None]
    ) * self.no_obj_embed_spatial[..., None, None]
    
    return maskmem_features, maskmem_pos_enc

def _prepare_memory_conditioned_features(
    self,
    frame_idx,
    is_init_cond_frame,
    current_vision_feats,
    output_dict,
    num_frames,
    track_in_reverse=False,
):
    """Condition current frame on memories from previous frames."""
    B = current_vision_feats[-1].size(1)
    C = self.hidden_dim
    H, W = feat_sizes[-1]
    device = current_vision_feats[-1].device
    
    if is_init_cond_frame:
        # First frame: no memory, add no_mem_embed
        pix_feat_with_mem = current_vision_feats[-1] + self.no_mem_embed
        return pix_feat_with_mem.permute(1, 2, 0).view(B, C, H, W)
    
    # Retrieve memories from conditioning frames
    to_cat_prompt, to_cat_prompt_pos_embed = [], []
    
    # Add conditioning frame outputs
    for t, out in output_dict["cond_frame_outputs"].items():
        t_pos = (frame_idx - t) * tpos_sign_mul
        
        # Spatial memory
        maskmem_features = out["maskmem_features"].cuda()
        to_cat_prompt.append(maskmem_features.flatten(2).permute(2, 0, 1))
        
        # Positional encoding
        maskmem_pos_enc = out["maskmem_pos_enc"][-1].cuda()
        maskmem_pos_enc = maskmem_pos_enc + self.maskmem_tpos_enc[self.num_maskmem - t_pos - 1]
        to_cat_prompt_pos_embed.append(maskmem_pos_enc.flatten(2).permute(2, 0, 1))
    
    # Add recent non-conditioning frames (last num_maskmem-1 frames)
    for t_pos in range(1, self.num_maskmem):
        prev_frame_idx = frame_idx - t_pos if not track_in_reverse else frame_idx + t_pos
        out = output_dict["non_cond_frame_outputs"].get(prev_frame_idx)
        if out is not None:
            # Add to memory
            ...
    
    # Collect object pointers
    for t, out in ptr_cond_outputs.items():
        obj_ptrs.append(out["obj_ptr"])
        obj_pos.append(self._get_tpos_enc([frame_idx - t], device))
    
    # Concatenate all prompts
    prompt = torch.cat(to_cat_prompt, dim=0)
    prompt_pos_embed = torch.cat(to_cat_prompt_pos_embed, dim=0)
    
    # Cross-attention: current frame attends to memories
    encoder_out = self.transformer.encoder(
        src=current_vision_feats,
        prompt=prompt,
        prompt_pos=prompt_pos_embed,
        num_obj_ptr_tokens=num_obj_ptr_tokens,
    )
    
    # Return memory-conditioned features
    pix_feat_with_mem = encoder_out["memory"].permute(1, 2, 0).view(B, C, H, W)
    return pix_feat_with_mem

class Sam3ImageOnVideoMultiGPU(Sam3Image):
    def forward_video_grounding_multigpu(
        self,
        frame_idx,
        num_frames,
        multigpu_buffer,  # Cache for detector outputs
        ...
    ):
        """
        Process frames in chunks across GPUs:
        - GPU 0 processes frame 0, 4, 8, ...
        - GPU 1 processes frame 1, 5, 9, ...
        - GPU 2 processes frame 2, 6, 10, ...
        - GPU 3 processes frame 3, 7, 11, ...
        
        Results are gathered via NCCL all-gather.
        """
        # Compute detection on local GPU's frame
        frame_idx_local = frame_idx_begin + self.rank
        out_local = self.forward_grounding(
            find_input=find_inputs[frame_idx_local],
            geometric_prompt=geometric_prompt,
        )
        
        # Gather results from all GPUs
        out_gathered = {k: self._gather_tensor(v) for k, v in out_local.items()}
        
        # Cache in multigpu_buffer for later access
        for rank in range(self.world_size):
            frame_idx_to_save = frame_idx_begin + rank
            multigpu_buffer[frame_idx_to_save] = out_gathered[rank]
        
        return multigpu_buffer