transformers 5.0.0__py3-none-any.whl → 5.0.0rc0__py3-none-any.whl

transformers/models/edgetam_video/modeling_edgetam_video.py

@@ -4,6 +4,7 @@
 #             the file from the modular. If any change should be done, please apply the change to the
 #                          modular_edgetam_video.py file directly. One of our CI enforces this.
 #                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# coding=utf-8
 # Copyright 2025 the HuggingFace Team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -22,7 +23,7 @@ import math
 from collections import OrderedDict
 from collections.abc import Callable, Iterator
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Optional, Union
 
 import numpy as np
 import torch
@@ -37,12 +38,12 @@ from ... import initialization as init
 from ...activations import ACT2FN
 from ...modeling_flash_attention_utils import FlashAttentionKwargs
 from ...modeling_layers import GradientCheckpointingLayer
-from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
+from ...modeling_outputs import BaseModelOutput
 from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from ...processing_utils import Unpack
 from ...pytorch_utils import compile_compatible_method_lru_cache
-from ...utils import ModelOutput, auto_docstring, can_return_tuple, logging
-from ...utils.generic import TransformersKwargs, is_flash_attention_requested
+from ...utils import ModelOutput, auto_docstring
+from ...utils.generic import TransformersKwargs
 from ..auto import AutoModel
 from .configuration_edgetam_video import (
     EdgeTamVideoConfig,
@@ -51,9 +52,6 @@ from .configuration_edgetam_video import (
 )
 
 
-logger = logging.get_logger(__name__)
-
-
 class EdgeTamVideoLayerNorm(nn.LayerNorm):
     r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
     The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height,
@@ -119,10 +117,16 @@ class EdgeTamVideoMemoryFuserCXBlock(GradientCheckpointingLayer):
 
 @dataclass
 @auto_docstring(custom_intro="Base class for the vision encoder's outputs.")
-class EdgeTamVideoVisionEncoderOutput(BaseModelOutputWithPooling):
+class EdgeTamVideoVisionEncoderOutput(ModelOutput):
     r"""
     last_hidden_state (`torch.FloatTensor` of shape `(batch_size, height, width, hidden_size)`):
         Sequence of hidden-states at the output of the last layer of the model.
+    fpn_hidden_states (`tuple(torch.FloatTensor)`):
+        Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape
+        `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck.
+    fpn_position_encoding (`tuple(torch.FloatTensor)`):
+        Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape
+        `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`.
     hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
         Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
         one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`. Hidden-states of the
@@ -131,16 +135,13 @@ class EdgeTamVideoVisionEncoderOutput(BaseModelOutputWithPooling):
         Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
         sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in
         the self-attention heads.
-    fpn_hidden_states (`tuple(torch.FloatTensor)`):
-        Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape
-        `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck.
-    fpn_position_encoding (`tuple(torch.FloatTensor)`):
-        Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape
-        `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`.
     """
 
-    fpn_hidden_states: torch.FloatTensor | None = None
-    fpn_position_encoding: torch.FloatTensor | None = None
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    fpn_hidden_states: Optional[torch.FloatTensor] = None
+    fpn_position_encoding: Optional[torch.FloatTensor] = None
+    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[tuple[torch.FloatTensor, ...]] = None
 
 
 class EdgeTamVideoVisionRotaryEmbedding(nn.Module):
@@ -149,19 +150,26 @@ class EdgeTamVideoVisionRotaryEmbedding(nn.Module):
     Supports 2D (axial) rotary embeddings for spatial dimensions.
     """
 
-    def __init__(self, config: EdgeTamVideoConfig, end_x: int | None = None, end_y: int | None = None):
+    def __init__(self, config: EdgeTamVideoConfig, end_x: Optional[int] = None, end_y: Optional[int] = None):
         super().__init__()
-        self.dim = config.memory_attention_hidden_size // (
+        dim = config.memory_attention_hidden_size // (
             config.memory_attention_downsample_rate * config.memory_attention_num_attention_heads
         )
         # Ensure even dimension for proper axial splitting
-        if self.dim % 4 != 0:
+        if dim % 4 != 0:
             raise ValueError("Dimension must be divisible by 4 for axial RoPE")
-        self.end_x, self.end_y = config.memory_attention_rope_feat_sizes if end_x is None else (end_x, end_y)
-        self.memory_attention_rope_theta = config.memory_attention_rope_theta
+        end_x, end_y = config.memory_attention_rope_feat_sizes if end_x is None else (end_x, end_y)
+        freqs = 1.0 / (config.memory_attention_rope_theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
 
+        # Generate 2D position indices for axial rotary embedding
+        flattened_indices = torch.arange(end_x * end_y, dtype=torch.long)
+        x_positions = flattened_indices % end_x
+        y_positions = torch.div(flattened_indices, end_x, rounding_mode="floor")
+        freqs_x = torch.outer(x_positions, freqs).float()
+        freqs_y = torch.outer(y_positions, freqs).float()
+        inv_freq = torch.cat([freqs_x, freqs_y], dim=-1)
+        inv_freq = inv_freq.repeat_interleave(2, dim=-1)
         # directly register the cos and sin embeddings as we have a fixed feature shape
-        inv_freq = self.create_inv_freq()
         self.register_buffer("rope_embeddings_cos", inv_freq.cos(), persistent=False)
         self.register_buffer("rope_embeddings_sin", inv_freq.sin(), persistent=False)
 
@@ -170,27 +178,13 @@ class EdgeTamVideoVisionRotaryEmbedding(nn.Module):
         # As the feature map size is fixed, we can just return the pre-computed embeddings.
         return self.rope_embeddings_cos, self.rope_embeddings_sin
 
-    def create_inv_freq(self):
-        freqs = 1.0 / (
-            self.memory_attention_rope_theta ** (torch.arange(0, self.dim, 4)[: (self.dim // 4)].float() / self.dim)
-        )
-        # Generate 2D position indices for axial rotary embedding
-        flattened_indices = torch.arange(self.end_x * self.end_y, dtype=torch.long)
-        x_positions = flattened_indices % self.end_x
-        y_positions = torch.div(flattened_indices, self.end_x, rounding_mode="floor")
-        freqs_x = torch.outer(x_positions, freqs).float()
-        freqs_y = torch.outer(y_positions, freqs).float()
-        inv_freq = torch.cat([freqs_x, freqs_y], dim=-1)
-        inv_freq = inv_freq.repeat_interleave(2, dim=-1)
-        return inv_freq
-
 
 def eager_attention_forward(
     module: nn.Module,
     query: torch.Tensor,
     key: torch.Tensor,
     value: torch.Tensor,
-    attention_mask: torch.Tensor | None,
+    attention_mask: Optional[torch.Tensor],
     scaling: float,
     dropout: float = 0.0,
     **kwargs,
@@ -234,7 +228,7 @@ class EdgeTamVideoAttention(nn.Module):
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attention_similarity: torch.Tensor | None = None,
+        attention_similarity: Optional[torch.Tensor] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> tuple[torch.Tensor, torch.Tensor]:
         # Input projections
@@ -249,15 +243,6 @@ class EdgeTamVideoAttention(nn.Module):
         if self.config._attn_implementation != "eager":
             attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
 
-        if is_flash_attention_requested(self.config) and attention_similarity is not None:
-            # Target guided masks are represented as float masks and are incompatible with Flash Attention
-            # Fallback to SDPA for this call only so the rest of the model can still benefit from FA
-            attention_interface = ALL_ATTENTION_FUNCTIONS["sdpa"]
-            logger.warning_once(
-                "Falling back to SDPA for target-guided attention because "
-                "Flash Attention does not support additive bias masks."
-            )
-
         attn_output, attn_weights = attention_interface(
             self,
             query,
@@ -622,7 +607,7 @@ class EdgeTamVideoPositionEmbeddingSine(nn.Module):
     """
 
     def __init__(
-        self, num_pos_feats: int = 64, temperature: int = 10000, normalize: bool = False, scale: float | None = None
+        self, num_pos_feats: int = 64, temperature: int = 10000, normalize: bool = False, scale: Optional[float] = None
     ):
         super().__init__()
         if scale is not None and normalize is False:
@@ -636,9 +621,9 @@ class EdgeTamVideoPositionEmbeddingSine(nn.Module):
     def forward(
         self,
         shape: torch.Size,
-        device: torch.device | str,
+        device: Union[torch.device, str],
         dtype: torch.dtype,
-        mask: Tensor | None = None,
+        mask: Optional[Tensor] = None,
     ) -> Tensor:
         if mask is None:
             mask = torch.zeros((shape[0], shape[2], shape[3]), device=device, dtype=torch.bool)
@@ -784,31 +769,6 @@ class EdgeTamVideoFeedForward(nn.Module):
         return hidden_states
 
 
-class EdgeTamVideoPositionalEmbedding(nn.Module):
-    def __init__(self, config: EdgeTamVideoPromptEncoderConfig):
-        super().__init__()
-        self.scale = config.scale
-        positional_embedding = self.scale * torch.randn((2, config.hidden_size // 2))
-        self.register_buffer("positional_embedding", positional_embedding)
-
-    def forward(self, input_coords, input_shape=None):
-        """Positionally encode points that are normalized to [0,1]."""
-        coordinates = input_coords.clone()
-
-        if input_shape is not None:
-            coordinates[:, :, :, 0] = coordinates[:, :, :, 0] / input_shape[1]
-            coordinates[:, :, :, 1] = coordinates[:, :, :, 1] / input_shape[0]
-        coordinates.to(torch.float32)
-
-        # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
-        coordinates = 2 * coordinates - 1
-        coordinates = coordinates.to(self.positional_embedding.dtype)
-        coordinates = coordinates @ self.positional_embedding
-        coordinates = 2 * np.pi * coordinates
-        # outputs d_1 x ... x d_n x channel shape
-        return torch.cat([torch.sin(coordinates), torch.cos(coordinates)], dim=-1)
-
-
 @auto_docstring
 class EdgeTamVideoPreTrainedModel(PreTrainedModel):
     config_class = EdgeTamVideoConfig
@@ -816,7 +776,7 @@ class EdgeTamVideoPreTrainedModel(PreTrainedModel):
     main_input_name = "pixel_values"
     input_modalities = "video"
     _supports_sdpa = True
-    _supports_flash_attn = True
+    _supports_flash_attn_2 = True
     _supports_attention_backend = True
 
     @torch.no_grad()
@@ -834,16 +794,6 @@ class EdgeTamVideoPreTrainedModel(PreTrainedModel):
         if isinstance(module, EdgeTamVideoMemoryFuserCXBlock):
             if module.scale is not None:
                 init.zeros_(module.scale)
-        elif isinstance(module, EdgeTamVideoVisionRotaryEmbedding):
-            inv_freq = module.create_inv_freq()
-            init.copy_(module.rope_embeddings_cos, inv_freq.cos())
-            init.copy_(module.rope_embeddings_sin, inv_freq.sin())
-        elif isinstance(module, EdgeTamVideoPositionalEmbedding):
-            init.normal_(module.positional_embedding, std=module.scale)
-        if isinstance(module, EdgeTamVideoVisionRotaryEmbedding):
-            inv_freq = module.create_inv_freq()
-            init.copy_(module.rope_embeddings_cos, inv_freq.cos())
-            init.copy_(module.rope_embeddings_sin, inv_freq.sin())
 
 
 class EdgeTamVideoInferenceCache:
@@ -851,8 +801,8 @@ class EdgeTamVideoInferenceCache:
 
     def __init__(
         self,
-        inference_device: torch.device | str = "cpu",
-        inference_state_device: torch.device | str = "cpu",
+        inference_device: Union[torch.device, str] = "cpu",
+        inference_state_device: Union[torch.device, str] = "cpu",
         max_vision_features_cache_size: int = 1,
     ):
         self.inference_device = inference_device
@@ -877,7 +827,7 @@ class EdgeTamVideoInferenceCache:
                 cached[key] = value
         self._vision_features[frame_idx] = cached
 
-    def get_vision_features(self, frame_idx: int) -> dict | None:
+    def get_vision_features(self, frame_idx: int) -> Optional[dict]:
         """Get cached vision features, automatically moved to inference device."""
         if frame_idx not in self._vision_features:
             return None
@@ -923,13 +873,13 @@ class EdgeTamVideoInferenceSession:
 
     def __init__(
         self,
-        video: torch.FloatTensor | None = None,
-        video_height: int | None = None,
-        video_width: int | None = None,
-        inference_device: torch.device | str = "cpu",
-        inference_state_device: torch.device | str = "cpu",
-        video_storage_device: torch.device | str = "cpu",
-        dtype: torch.dtype | str = "float32",
+        video: Optional[torch.FloatTensor] = None,
+        video_height: Optional[int] = None,
+        video_width: Optional[int] = None,
+        inference_device: Union[torch.device, str] = "cpu",
+        inference_state_device: Union[torch.device, str] = "cpu",
+        video_storage_device: Union[torch.device, str] = "cpu",
+        dtype: Union[torch.dtype, str] = "float32",
         max_vision_features_cache_size: int = 1,
     ):
         # store as a dictionary to avoid double memory allocation with torch.cat when adding new frames
@@ -969,7 +919,7 @@ class EdgeTamVideoInferenceSession:
         self.obj_with_new_inputs = []
 
     @property
-    def num_frames(self) -> int | None:
+    def num_frames(self) -> Optional[int]:
         return len(self.processed_frames) if self.processed_frames is not None else None
 
     # Object management
@@ -1009,7 +959,7 @@ class EdgeTamVideoInferenceSession:
         device_inputs = {}
         for key, value in inputs.items():
             if isinstance(value, torch.Tensor):
-                device_inputs[key] = value.to(self.inference_device, non_blocking=False)
+                device_inputs[key] = value.to(self.inference_device, non_blocking=True)
             else:
                 device_inputs[key] = value
         self.point_inputs_per_obj[obj_idx][frame_idx] = device_inputs
@@ -1033,8 +983,8 @@ class EdgeTamVideoInferenceSession:
         self,
         obj_idx: int,
         frame_idx: int,
-        output_key: str | None = None,
-        output_value: torch.Tensor | dict | None = None,
+        output_key: Optional[str] = None,
+        output_value: Optional[Union[torch.Tensor, dict]] = None,
         is_conditioning_frame: bool = True,
     ):
         """
@@ -1093,7 +1043,7 @@ class EdgeTamVideoInferenceSession:
         return value
 
     # Video frame management
-    def add_new_frame(self, pixel_values: torch.Tensor, frame_idx: int | None = None) -> int:
+    def add_new_frame(self, pixel_values: torch.Tensor, frame_idx: Optional[int] = None) -> int:
         """Add new frame with automatic device placement."""
         pixel_values = pixel_values.to(self.video_storage_device, dtype=self.dtype, non_blocking=True)
         if pixel_values.dim() == 4:
@@ -1176,7 +1126,7 @@ class EdgeTamVideoMemoryAttentionLayer(nn.Module):
         keys: Tensor,
         key_point_embedding: Tensor,
         rope_position_embeddings: tuple[Tensor, Tensor],
-        rope_position_embeddings_k: tuple[Tensor, Tensor] | None = None,
+        rope_position_embeddings_k: Optional[tuple[Tensor, Tensor]] = None,
         num_k_exclude_rope: int = 0,
         rope_k_repeat: int = 0,
     ) -> torch.Tensor:
@@ -1220,8 +1170,8 @@ class EdgeTamVideoMemoryAttention(nn.Module):
         self,
         current_vision_features: torch.Tensor,
         memory: torch.Tensor,
-        current_vision_position_embeddings: Tensor | None = None,
-        memory_posision_embeddings: Tensor | None = None,
+        current_vision_position_embeddings: Optional[Tensor] = None,
+        memory_posision_embeddings: Optional[Tensor] = None,
         num_object_pointer_tokens: int = 0,
         num_spatial_memory_tokens: int = -1,
     ):
@@ -1307,7 +1257,7 @@ class EdgeTamVideoPerceiverAttention(nn.Module):
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        positional_encoding: torch.Tensor | None = None,
+        positional_encoding: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
         # Project queries, keys, and values
@@ -1372,7 +1322,7 @@ class EdgeTamVideoPerceiverEncoderLayer(nn.Module):
         self,
         latents: torch.Tensor,
         input_features: torch.Tensor,
-        positional_encoding: torch.Tensor | None = None,
+        positional_encoding: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         # Cross attention with layer norms
         normalized_latents = self.layer_norm_latents(latents)
@@ -1458,8 +1408,8 @@ class EdgeTamVideoPerceiverResampler(nn.Module):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        positional_encoding: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        positional_encoding: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         output_latents = []
         output_positional_encodings = []
 
@@ -1484,8 +1434,8 @@ class EdgeTamVideoPerceiverResampler(nn.Module):
     def _forward_1d(
         self,
         hidden_states: torch.Tensor,
-        positional_encoding: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        positional_encoding: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         batch_size = hidden_states.shape[0]
 
         latents = self.latents_1d.unsqueeze(0).expand(batch_size, -1, -1)
@@ -1565,16 +1515,16 @@ class EdgeTamVideoImageSegmentationOutput(ModelOutput):
         A tensor representing the object pointer, used for tracking in videos. Only used for EdgeTamVideoModel.
     """
 
-    iou_scores: torch.FloatTensor | None = None
-    pred_masks: torch.FloatTensor | None = None
-    object_score_logits: torch.FloatTensor | None = None
+    iou_scores: Optional[torch.FloatTensor] = None
+    pred_masks: Optional[torch.FloatTensor] = None
+    object_score_logits: Optional[torch.FloatTensor] = None
     image_embeddings: tuple[torch.FloatTensor, ...] = None
-    vision_hidden_states: tuple[torch.FloatTensor, ...] | None = None
-    vision_attentions: tuple[torch.FloatTensor, ...] | None = None
-    mask_decoder_attentions: tuple[torch.FloatTensor, ...] | None = None
+    vision_hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+    vision_attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+    mask_decoder_attentions: Optional[tuple[torch.FloatTensor, ...]] = None
 
-    high_res_masks: torch.FloatTensor | None = None
-    object_pointer: torch.FloatTensor | None = None
+    high_res_masks: Optional[torch.FloatTensor] = None
+    object_pointer: Optional[torch.FloatTensor] = None
 
 
 @dataclass
@@ -1591,10 +1541,35 @@ class EdgeTamVideoSegmentationOutput(ModelOutput):
         The frame index of the video.
     """
 
-    object_ids: list[int] | None = None
-    pred_masks: torch.FloatTensor | None = None
-    object_score_logits: torch.FloatTensor | None = None
-    frame_idx: int | None = None
+    object_ids: Optional[list[int]] = None
+    pred_masks: Optional[torch.FloatTensor] = None
+    object_score_logits: Optional[torch.FloatTensor] = None
+    frame_idx: Optional[int] = None
+
+
+class EdgeTamVideoPositionalEmbedding(nn.Module):
+    def __init__(self, config: EdgeTamVideoPromptEncoderConfig):
+        super().__init__()
+        self.scale = config.scale
+        positional_embedding = self.scale * torch.randn((2, config.hidden_size // 2))
+        self.register_buffer("positional_embedding", positional_embedding)
+
+    def forward(self, input_coords, input_shape=None):
+        """Positionally encode points that are normalized to [0,1]."""
+        coordinates = input_coords.clone()
+
+        if input_shape is not None:
+            coordinates[:, :, :, 0] = coordinates[:, :, :, 0] / input_shape[1]
+            coordinates[:, :, :, 1] = coordinates[:, :, :, 1] / input_shape[0]
+        coordinates.to(torch.float32)
+
+        # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
+        coordinates = 2 * coordinates - 1
+        coordinates = coordinates.to(self.positional_embedding.dtype)
+        coordinates = coordinates @ self.positional_embedding
+        coordinates = 2 * np.pi * coordinates
+        # outputs d_1 x ... x d_n x channel shape
+        return torch.cat([torch.sin(coordinates), torch.cos(coordinates)], dim=-1)
 
 
 class EdgeTamVideoMaskEmbedding(nn.Module):
@@ -1678,10 +1653,10 @@ class EdgeTamVideoPromptEncoder(nn.Module):
 
     def forward(
         self,
-        input_points: tuple[torch.Tensor, torch.Tensor] | None,
-        input_labels: torch.Tensor | None,
-        input_boxes: torch.Tensor | None,
-        input_masks: torch.Tensor | None,
+        input_points: Optional[tuple[torch.Tensor, torch.Tensor]],
+        input_labels: Optional[torch.Tensor],
+        input_boxes: Optional[torch.Tensor],
+        input_masks: Optional[torch.Tensor],
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Embeds different types of prompts, returning both sparse and dense embeddings.
@@ -1741,7 +1716,7 @@ class EdgeTamVideoTwoWayTransformer(nn.Module):
         attention_similarity: Tensor,
         target_embedding=None,
         **kwargs: Unpack[TransformersKwargs],
-    ) -> tuple | BaseModelOutput:
+    ) -> Union[tuple, BaseModelOutput]:
         if image_embeddings is None:
             raise ValueError("You have to specify an image_embedding")
 
@@ -1826,8 +1801,8 @@ class EdgeTamVideoMaskDecoder(nn.Module):
         dense_prompt_embeddings: torch.Tensor,
         multimask_output: bool,
         high_resolution_features: list[torch.Tensor],
-        attention_similarity: torch.Tensor | None = None,
-        target_embedding: torch.Tensor | None = None,
+        attention_similarity: Optional[torch.Tensor] = None,
+        target_embedding: Optional[torch.Tensor] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         """
@@ -2000,8 +1975,12 @@ def get_1d_sine_pe(pos_inds, dim, temperature=10000):
 class EdgeTamVideoModel(EdgeTamVideoPreTrainedModel):
     input_modalities = ("video", "text")
     _can_record_outputs = {"mask_decoder_attentions": OutputRecorder(EdgeTamVideoTwoWayAttentionBlock, index=2)}
-    _tied_weights_keys = {}
     _keys_to_ignore_on_load_unexpected = []
+    _tied_weights_keys = {
+        "prompt_encoder.shared_embedding.positional_embedding": "shared_image_embedding.positional_embedding"
+    }
+    # need to be ignored, as it's a buffer and will not be correctly detected as tied weight
+    _keys_to_ignore_on_load_missing = ["prompt_encoder.shared_embedding.positional_embedding"]
 
     def __init__(self, config: EdgeTamVideoConfig):
         super().__init__(config)
@@ -2084,8 +2063,7 @@ class EdgeTamVideoModel(EdgeTamVideoPreTrainedModel):
                 Input pixel values
         """
         batch_size = pixel_values.shape[0]
-        image_outputs = self.get_image_features(pixel_values, return_dict=True, **kwargs)
-        feature_maps = image_outputs.fpn_hidden_states
+        feature_maps, _, _, _ = self.get_image_features(pixel_values, **kwargs)
 
         # add no memory embedding to the last feature map
         feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding
@@ -2101,10 +2079,10 @@ class EdgeTamVideoModel(EdgeTamVideoPreTrainedModel):
     @torch.no_grad()
     def get_prompt_embeddings(
         self,
-        input_points: torch.FloatTensor | None = None,
-        input_labels: torch.LongTensor | None = None,
-        input_boxes: torch.FloatTensor | None = None,
-        input_masks: torch.LongTensor | None = None,
+        input_points: Optional[torch.FloatTensor] = None,
+        input_labels: Optional[torch.LongTensor] = None,
+        input_boxes: Optional[torch.FloatTensor] = None,
+        input_masks: Optional[torch.LongTensor] = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         r"""
         Returns the prompt embeddings by passing the input points, labels, boxes and masks through the prompt encoder.
@@ -2136,10 +2114,9 @@ class EdgeTamVideoModel(EdgeTamVideoPreTrainedModel):
     def forward(
         self,
         inference_session: EdgeTamVideoInferenceSession,
-        frame_idx: int | None = None,
-        frame: torch.Tensor | None = None,
+        frame_idx: Optional[int] = None,
+        frame: Optional[torch.Tensor] = None,
         reverse: bool = False,
-        **kwargs,
     ) -> EdgeTamVideoSegmentationOutput:
         r"""
         inference_session (`EdgeTamVideoInferenceSession`):
@@ -2230,18 +2207,34 @@ class EdgeTamVideoModel(EdgeTamVideoPreTrainedModel):
             frame_idx=frame_idx,
         )
 
-    @can_return_tuple
-    @auto_docstring
     def get_image_features(
         self,
         pixel_values: torch.FloatTensor,
         **kwargs: Unpack[TransformersKwargs],
-    ) -> tuple | EdgeTamVideoVisionEncoderOutput:
+    ) -> tuple[
+        list[torch.Tensor],
+        list[torch.Tensor],
+        Optional[tuple[torch.FloatTensor, ...]],
+        Optional[tuple[torch.FloatTensor, ...]],
+    ]:
         r"""
-        pixel_values (`torch.FloatTensor`):
-            Input pixel values of shape `(batch_size, num_channels, height, width)`.
+        Extract and preprocess image features using the vision encoder.
+
+        Args:
+            pixel_values (`torch.FloatTensor`):
+                Input pixel values of shape `(batch_size, num_channels, height, width)`.
+
+        Returns:
+            `tuple`: A tuple containing:
+                - feature_maps (`list[torch.Tensor]`): List of feature maps from different levels.
+                - feature_maps_position_embeddings (`list[torch.Tensor]`): List of positional embeddings for each feature level.
+                - vision_hidden_states (`tuple[torch.FloatTensor]`, *optional*): Hidden states from the vision encoder.
+                - vision_attentions (`tuple[torch.FloatTensor]`, *optional*): Attention weights from the vision encoder.
         """
-        vision_outputs: EdgeTamVideoVisionEncoderOutput = self.vision_encoder(pixel_values, return_dict=True, **kwargs)
+        vision_outputs: EdgeTamVideoVisionEncoderOutput = self.vision_encoder(
+            pixel_values,
+            **kwargs,
+        )
 
         feature_maps = vision_outputs.fpn_hidden_states
         feature_maps_position_embeddings = vision_outputs.fpn_position_encoding
@@ -2258,10 +2251,8 @@ class EdgeTamVideoModel(EdgeTamVideoPreTrainedModel):
             feature_map_position_embedding.flatten(2).permute(2, 0, 1)
             for feature_map_position_embedding in feature_maps_position_embeddings
         ]
-        vision_outputs.fpn_hidden_states = feature_maps
-        vision_outputs.fpn_position_encoding = feature_maps_position_embeddings
 
-        return vision_outputs
+        return feature_maps, feature_maps_position_embeddings, vision_outputs.hidden_states, vision_outputs.attentions
 
     def _prepare_vision_features(
         self,
@@ -2278,9 +2269,7 @@ class EdgeTamVideoModel(EdgeTamVideoPreTrainedModel):
         else:
             # Compute features using image encoder
             image_batch = inference_session.get_frame(frame_idx).unsqueeze(0)  # Add batch dimension
-            image_outputs = self.get_image_features(image_batch, return_dict=True)
-            vision_feats = image_outputs.fpn_hidden_states
-            vision_pos_embeds = image_outputs.fpn_position_embeddings
+            vision_feats, vision_pos_embeds, _, _ = self.get_image_features(image_batch)
             # Cache features
             inference_session.cache.cache_vision_features(
                 frame_idx, {"vision_feats": vision_feats, "vision_pos_embeds": vision_pos_embeds}
@@ -2295,15 +2284,15 @@ class EdgeTamVideoModel(EdgeTamVideoPreTrainedModel):
 
     def _single_frame_forward(
         self,
-        pixel_values: torch.FloatTensor | None = None,
-        input_points: torch.FloatTensor | None = None,
-        input_labels: torch.LongTensor | None = None,
-        input_boxes: torch.FloatTensor | None = None,
-        input_masks: torch.LongTensor | None = None,
-        image_embeddings: torch.FloatTensor | None = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        input_points: Optional[torch.FloatTensor] = None,
+        input_labels: Optional[torch.LongTensor] = None,
+        input_boxes: Optional[torch.FloatTensor] = None,
+        input_masks: Optional[torch.LongTensor] = None,
+        image_embeddings: Optional[torch.FloatTensor] = None,
         multimask_output: bool = True,
-        attention_similarity: torch.FloatTensor | None = None,
-        target_embedding: torch.FloatTensor | None = None,
+        attention_similarity: Optional[torch.FloatTensor] = None,
+        target_embedding: Optional[torch.FloatTensor] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> EdgeTamVideoImageSegmentationOutput:
         """
@@ -2385,10 +2374,10 @@ class EdgeTamVideoModel(EdgeTamVideoPreTrainedModel):
         vision_hidden_states = None
 
         if pixel_values is not None:
-            image_outputs = self.get_image_features(pixel_values, return_dict=True, **kwargs)
-            feature_maps = image_outputs.fpn_hidden_states
-            vision_hidden_states = image_outputs.hidden_states
-            vision_attentions = image_outputs.attentions
+            feature_maps, _, vision_hidden_states, vision_attentions = self.get_image_features(
+                pixel_values,
+                **kwargs,
+            )
 
             # add no memory embedding to the last feature map
             feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding
@@ -2888,7 +2877,7 @@ class EdgeTamVideoModel(EdgeTamVideoPreTrainedModel):
         )
         return conditioned_feature_map
 
-    def _use_multimask(self, is_init_cond_frame: bool, point_inputs: dict | None) -> bool:
+    def _use_multimask(self, is_init_cond_frame: bool, point_inputs: Optional[dict]) -> bool:
         """Whether to use multimask output in the SAM head."""
         num_pts = 0 if point_inputs is None else point_inputs["point_labels"].size(2)
         multimask_output = (
@@ -2905,11 +2894,11 @@ class EdgeTamVideoModel(EdgeTamVideoPreTrainedModel):
         obj_idx: int,
         batch_size: int,
         is_init_cond_frame: bool,
-        point_inputs: torch.Tensor | None,
-        mask_inputs: torch.Tensor | None,
+        point_inputs: Optional[torch.Tensor],
+        mask_inputs: Optional[torch.Tensor],
         reverse: bool,
         run_mem_encoder: bool,
-        prev_sam_mask_logits: torch.Tensor | None = None,
+        prev_sam_mask_logits: Optional[torch.Tensor] = None,
         streaming: bool = False,
     ) -> dict[str, Any]:
         """
@@ -3073,8 +3062,8 @@ class EdgeTamVideoModel(EdgeTamVideoPreTrainedModel):
     def propagate_in_video_iterator(
         self,
         inference_session: EdgeTamVideoInferenceSession,
-        start_frame_idx: int | None = None,
-        max_frame_num_to_track: int | None = None,
+        start_frame_idx: Optional[int] = None,
+        max_frame_num_to_track: Optional[int] = None,
         reverse: bool = False,
         show_progress_bar: bool = False,
     ) -> Iterator[EdgeTamVideoSegmentationOutput]: