Loading embeddings to speed up video predictions

[jfoody-centaur]

For the SAM2 image predictor, I created a custom initialization of the image predictor which can save/load the image embeddings to speed up the process when those images are called for prediction again.

I'm interested in doing something similar for the SAM2 video predictor where I preload the embeddings for all frames. Is this (or something similar) possible for the video predictor as the embeddings/memory attention seem to differ from the image predictor.

I'm also having trouble getting the model to initialize correctly with a custom initialization using something like this:

class MySamVideo(SAM2VideoPredictor):
    def __init__(self, model, **kwargs):
        # Extract required attributes from `model` if available
        extracted_args = {
            "image_encoder": getattr(model, "image_encoder", None),
            "memory_attention": getattr(model, "memory_attention", None),
            "memory_encoder": getattr(model, "memory_encoder", None)
        }

        super().__init__(**extracted_args, **kwargs)
        self.model = model

    def save_image_embedding(self, path):
        # Save existing embeddings

    def load_image_embedding(self, path):
        # Check if embedding exists & load

# called like this
sam = build_sam2_video_predictor(model_cfg, model_path, device=device)
sam.to(device=device)
model = MySamVideo(sam)