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

transformers/models/sam3/modeling_sam3.py

@@ -14,9 +14,8 @@
 # limitations under the License.
 
 
-import collections.abc
 import math
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
 from dataclasses import dataclass
 from typing import Optional, Union
 
@@ -40,7 +39,7 @@ from ...modeling_outputs import (
 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 auto_docstring
+from ...utils import auto_docstring, logging
 from ...utils.generic import TransformersKwargs, check_model_inputs
 from ..auto import AutoModel
 from .configuration_sam3 import (
@@ -54,6 +53,9 @@ from .configuration_sam3 import (
 )
 
 
+logger = logging.get_logger(__name__)
+
+
 @dataclass
 @auto_docstring
 class Sam3VisionEncoderOutput(ModelOutput):
@@ -123,8 +125,8 @@ class Sam3DETRDecoderOutput(ModelOutput):
         Decoder hidden states from all layers.
     reference_boxes (`torch.FloatTensor` of shape `(num_layers, batch_size, num_queries, 4)`):
         Predicted reference boxes from all decoder layers in (cx, cy, w, h) format.
-    presence_logits (`torch.FloatTensor` of shape `(num_layers, batch_size)`, *optional*):
-        Presence logits from all decoder layers (None if using instance queries).
+    presence_logits (`torch.FloatTensor` of shape `(num_layers, batch_size, 1)`):
+        Presence logits from all decoder layers indicating object presence confidence.
     hidden_states (`tuple[torch.FloatTensor]`, *optional*):
         Tuple of hidden states from all decoder layers.
     attentions (`tuple[torch.FloatTensor]`, *optional*):
@@ -133,7 +135,7 @@ class Sam3DETRDecoderOutput(ModelOutput):
 
     intermediate_hidden_states: torch.FloatTensor = None
     reference_boxes: torch.FloatTensor = None
-    presence_logits: Optional[torch.FloatTensor] = None
+    presence_logits: torch.FloatTensor = None
     hidden_states: Optional[tuple[torch.FloatTensor]] = None
     attentions: Optional[tuple[torch.FloatTensor]] = None
 
@@ -372,6 +374,19 @@ class Sam3Attention(nn.Module):
         if self.config._attn_implementation != "eager":
             attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
 
+        if (
+            "flash" in self.config._attn_implementation
+            and attention_mask is not None
+            and attention_mask.dtype != torch.bool
+        ):
+            # Relative position bias tensors 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(
+                "Sam3Attention: falling back to SDPA for relative-position cross-attention because "
+                "Flash Attention does not support additive bias masks."
+            )
+
         attn_output, attn_weights = attention_interface(
             self,
             query,
@@ -402,6 +417,10 @@ class Sam3ViTRotaryEmbedding(nn.Module):
         # Ensure even dimension for proper axial splitting
         if dim % 4 != 0:
             raise ValueError("Dimension must be divisible by 4 for axial RoPE")
+        self.end_x, self.end_y = end_x, end_y
+        self.dim = dim
+        self.rope_theta = config.rope_theta
+        self.scale = scale
         freqs = 1.0 / (config.rope_theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
 
         flattened_indices = torch.arange(end_x * end_y, dtype=torch.long)
@@ -531,8 +550,8 @@ class Sam3ViTPatchEmbeddings(nn.Module):
         image_size, patch_size = config.pretrain_image_size, config.patch_size
         num_channels, hidden_size = config.num_channels, config.hidden_size
 
-        image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
-        patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
+        image_size = image_size if isinstance(image_size, Iterable) else (image_size, image_size)
+        patch_size = patch_size if isinstance(patch_size, Iterable) else (patch_size, patch_size)
         num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
         self.image_size = image_size
         self.patch_size = patch_size
@@ -542,7 +561,7 @@ class Sam3ViTPatchEmbeddings(nn.Module):
         self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size, bias=False)
 
     def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
-        embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2)
+        embeddings = self.projection(pixel_values.to(self.projection.weight.dtype)).flatten(2).transpose(1, 2)
         return embeddings
 
 
@@ -761,6 +780,19 @@ class Sam3PreTrainedModel(PreTrainedModel):
         super()._init_weights(module)
         if isinstance(module, Sam3ViTEmbeddings):
             init.normal_(module.position_embeddings, mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, Sam3ViTRotaryEmbedding):
+            end_x, end_y = module.end_x, module.end_y
+            dim = module.dim
+            freqs = 1.0 / (module.rope_theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
+            flattened_indices = torch.arange(end_x * end_y, dtype=torch.long)
+            x_positions = (flattened_indices % end_x) * module.scale
+            y_positions = torch.div(flattened_indices, end_x, rounding_mode="floor") * module.scale
+            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)
+            init.copy_(module.rope_embeddings_cos, inv_freq.cos())
+            init.copy_(module.rope_embeddings_sin, inv_freq.sin())
 
 
 @auto_docstring
@@ -938,6 +970,7 @@ class Sam3FPNLayer(nn.Module):
         self.proj2 = nn.Conv2d(in_channels=fpn_dim, out_channels=fpn_dim, kernel_size=3, padding=1)
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.to(self.proj1.weight.dtype)
         for layer in self.scale_layers:
             hidden_states = layer(hidden_states)
 
@@ -1253,7 +1286,7 @@ class Sam3DetrEncoderLayer(nn.Module):
         vision_feats: Tensor,
         prompt_feats: Tensor,
         vision_pos_encoding: Tensor,
-        prompt_mask: Tensor,
+        prompt_cross_attn_mask: Optional[Tensor] = None,
         **kwargs: Unpack[TransformersKwargs],
     ):
         """
@@ -1263,7 +1296,7 @@ class Sam3DetrEncoderLayer(nn.Module):
             vision_feats: Vision features [batch_size, vision_len, hidden_size] (main hidden states)
             prompt_feats: Text prompt features [batch_size, text_len, hidden_size]
             vision_pos_encoding: Position encoding for vision [batch_size, vision_len, hidden_size]
-            prompt_mask: Padding mask for prompts [batch_size, text_len] where True=valid, False=padding
+            prompt_cross_attn_mask: Cross-attention mask for prompt features
 
         Returns:
             Updated vision features [batch_size, vision_len, hidden_size]
@@ -1284,15 +1317,6 @@ class Sam3DetrEncoderLayer(nn.Module):
         residual = hidden_states
         hidden_states = self.layer_norm2(hidden_states)
 
-        prompt_cross_attn_mask = None
-        if prompt_mask is not None:
-            prompt_cross_attn_mask = create_bidirectional_mask(
-                config=self.config,
-                input_embeds=hidden_states,
-                attention_mask=prompt_mask,
-                encoder_hidden_states=prompt_feats,
-            )
-
         hidden_states, _ = self.cross_attn(
             query=hidden_states,
             key=prompt_feats,
@@ -1331,6 +1355,8 @@ class Sam3DetrEncoder(Sam3PreTrainedModel):
 
         self.layers = nn.ModuleList([Sam3DetrEncoderLayer(config) for _ in range(config.num_layers)])
 
+        self.post_init()
+
     def _prepare_multilevel_features(
         self,
         vision_features: list[torch.Tensor],
@@ -1412,13 +1438,22 @@ class Sam3DetrEncoder(Sam3PreTrainedModel):
             spatial_shapes,
         ) = self._prepare_multilevel_features(vision_features, vision_pos_embeds)
 
+        prompt_cross_attn_mask = None
+        if text_mask is not None:
+            prompt_cross_attn_mask = create_bidirectional_mask(
+                config=self.config,
+                input_embeds=features_flattened,
+                attention_mask=text_mask,
+                encoder_hidden_states=text_features,
+            )
+
         hidden_states = features_flattened
         for layer in self.layers:
             hidden_states = layer(
                 hidden_states,
                 prompt_feats=text_features,
                 vision_pos_encoding=pos_embeds_flattened,
-                prompt_mask=text_mask,
+                prompt_cross_attn_mask=prompt_cross_attn_mask,
                 **kwargs,
             )
         return Sam3DETREncoderOutput(
@@ -1484,31 +1519,27 @@ class Sam3DetrDecoderLayer(nn.Module):
         text_features: torch.Tensor,
         vision_features: torch.Tensor,
         vision_pos_encoding: torch.Tensor,
-        text_mask: Optional[torch.Tensor] = None,
+        text_cross_attn_mask: Optional[torch.Tensor] = None,
         vision_cross_attn_mask: Optional[torch.Tensor] = None,
-        presence_token: Optional[torch.Tensor] = None,
         **kwargs: Unpack[TransformersKwargs],
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+    ) -> torch.Tensor:
         """
         Forward pass for decoder layer.
 
         Args:
-            hidden_states: Query features [batch_size, num_queries, hidden_size]
+            hidden_states: Query features [batch_size, num_queries + 1, hidden_size] (includes presence token at position 0)
             query_pos: Query position embeddings [batch_size, num_queries, hidden_size]
             text_features: Text features [batch_size, seq_len, hidden_size]
             vision_features: Vision features [batch_size, height*width, hidden_size]
             vision_pos_encoding: Vision position encoding [batch_size, height*width, hidden_size]
-            text_mask: Text padding mask [batch_size, seq_len] where True=valid, False=padding
-            vision_cross_attn_mask: Vision cross-attention mask [batch_size, num_heads, num_queries, height*width]
-            presence_token: Optional presence token [batch_size, 1, hidden_size]
+            text_cross_attn_mask: Text cross-attention mask
+            vision_cross_attn_mask: Vision cross-attention mask, already expanded for presence token
 
         Returns:
-            Tuple of (updated hidden states, updated presence token)
+            Updated hidden states (including presence token at position 0)
         """
-        # Concatenate presence token if provided
-        if presence_token is not None:
-            hidden_states = torch.cat([presence_token, hidden_states], dim=1)
-            query_pos = torch.cat([torch.zeros_like(presence_token), query_pos], dim=1)
+        # Prepend zeros to query_pos for presence token
+        query_pos = F.pad(query_pos, (0, 0, 1, 0), mode="constant", value=0)
 
         # Self-attention with query position encoding
         residual = hidden_states
@@ -1527,15 +1558,6 @@ class Sam3DetrDecoderLayer(nn.Module):
         residual = hidden_states
         query_with_pos = hidden_states + query_pos
 
-        text_cross_attn_mask = None
-        if text_mask is not None:
-            text_cross_attn_mask = create_bidirectional_mask(
-                config=self.config,
-                input_embeds=hidden_states,
-                attention_mask=text_mask,
-                encoder_hidden_states=text_features,
-            )
-
         attn_output, _ = self.text_cross_attn(
             query=query_with_pos,
             key=text_features,
@@ -1546,20 +1568,6 @@ class Sam3DetrDecoderLayer(nn.Module):
         hidden_states = residual + self.text_cross_attn_dropout(attn_output)
         hidden_states = self.text_cross_attn_layer_norm(hidden_states)
 
-        # Expand vision cross-attention mask for presence token if needed
-        combined_vision_mask = vision_cross_attn_mask
-        if presence_token is not None and combined_vision_mask is not None:
-            batch_size, num_heads = combined_vision_mask.shape[:2]
-            presence_mask = torch.zeros(
-                batch_size,
-                num_heads,
-                1,
-                combined_vision_mask.shape[-1],
-                device=combined_vision_mask.device,
-                dtype=combined_vision_mask.dtype,
-            )
-            combined_vision_mask = torch.cat([presence_mask, combined_vision_mask], dim=2)
-
         # Vision cross-attention: queries attend to vision features (with RPB)
         residual = hidden_states
         query_with_pos = hidden_states + query_pos
@@ -1568,7 +1576,7 @@ class Sam3DetrDecoderLayer(nn.Module):
             query=query_with_pos,
             key=key_with_pos,
             value=vision_features,
-            attention_mask=combined_vision_mask,
+            attention_mask=vision_cross_attn_mask,
             **kwargs,
         )
         hidden_states = residual + self.vision_cross_attn_dropout(attn_output)
@@ -1580,13 +1588,7 @@ class Sam3DetrDecoderLayer(nn.Module):
         hidden_states = residual + self.mlp_dropout(hidden_states)
         hidden_states = self.mlp_layer_norm(hidden_states)
 
-        # Extract presence token if it was added
-        presence_token_out = None
-        if presence_token is not None:
-            presence_token_out = hidden_states[:, :1]
-            hidden_states = hidden_states[:, 1:]
-
-        return hidden_states, presence_token_out
+        return hidden_states
 
 
 class Sam3DetrDecoder(Sam3PreTrainedModel):
@@ -1634,6 +1636,8 @@ class Sam3DetrDecoder(Sam3PreTrainedModel):
 
         self.position_encoding = Sam3SinePositionEmbedding(num_pos_feats=config.hidden_size // 2, normalize=False)
 
+        self.post_init()
+
     @compile_compatible_method_lru_cache(maxsize=1)
     def _get_coords(
         self, height: torch.Tensor, width: torch.Tensor, dtype: torch.dtype, device: torch.device
@@ -1715,11 +1719,23 @@ class Sam3DetrDecoder(Sam3PreTrainedModel):
         """
         batch_size = vision_features.shape[0]
 
-        hidden_states = self.query_embed.weight.unsqueeze(0).expand(batch_size, -1, -1)
+        query_embeds = self.query_embed.weight.unsqueeze(0).expand(batch_size, -1, -1)
         reference_boxes = self.reference_points.weight.unsqueeze(0).expand(batch_size, -1, -1)
         reference_boxes = reference_boxes.sigmoid()
         presence_token = self.presence_token.weight.unsqueeze(0).expand(batch_size, -1, -1)
 
+        # Concatenate presence token with query embeddings
+        hidden_states = torch.cat([presence_token, query_embeds], dim=1)
+
+        text_cross_attn_mask = None
+        if text_mask is not None:
+            text_cross_attn_mask = create_bidirectional_mask(
+                config=self.config,
+                input_embeds=hidden_states,
+                attention_mask=text_mask,
+                encoder_hidden_states=text_features,
+            )
+
         intermediate_outputs = []
         intermediate_boxes = [reference_boxes]
         intermediate_presence_logits = []
@@ -1734,43 +1750,45 @@ class Sam3DetrDecoder(Sam3PreTrainedModel):
             vision_cross_attn_mask = None
             if spatial_shapes is not None and spatial_shapes.shape[0] == 1:
                 spatial_shape = (spatial_shapes[0, 0], spatial_shapes[0, 1])
-                vision_cross_attn_mask = self._get_rpb_matrix(reference_boxes, spatial_shape)
+                rpb_matrix = self._get_rpb_matrix(reference_boxes, spatial_shape)
+                # Prepend zeros row for presence token (it attends to all vision tokens equally)
+                vision_cross_attn_mask = F.pad(rpb_matrix, (0, 0, 1, 0), mode="constant", value=0)
 
-            hidden_states, presence_token = layer(
+            hidden_states = layer(
                 hidden_states,
                 query_pos=query_pos,
                 text_features=text_features,
                 vision_features=vision_features,
                 vision_pos_encoding=vision_pos_encoding,
-                text_mask=text_mask,
+                text_cross_attn_mask=text_cross_attn_mask,
                 vision_cross_attn_mask=vision_cross_attn_mask,
-                presence_token=presence_token,
                 **kwargs,
             )
 
+            # Extract query hidden states (without presence token) for box refinement
+            query_hidden_states = hidden_states[:, 1:]
+
             # Box refinement: predict delta and update reference boxes
             reference_boxes_before_sigmoid = inverse_sigmoid(reference_boxes)
-            delta_boxes = self.box_head(self.output_layer_norm(hidden_states))
+            delta_boxes = self.box_head(self.output_layer_norm(query_hidden_states))
             new_reference_boxes = (delta_boxes + reference_boxes_before_sigmoid).sigmoid()
             reference_boxes = new_reference_boxes.detach()
 
-            intermediate_outputs.append(self.output_layer_norm(hidden_states))
+            intermediate_outputs.append(self.output_layer_norm(query_hidden_states))
             intermediate_boxes.append(new_reference_boxes)
 
             # Process presence token
-            if presence_token is not None:
-                presence_logits = self.presence_head(self.presence_layer_norm(presence_token)).squeeze(-1)
-                presence_logits = presence_logits.clamp(
-                    min=-self.clamp_presence_logit_max_val, max=self.clamp_presence_logit_max_val
-                )
-                intermediate_presence_logits.append(presence_logits)
+            presence_hidden = hidden_states[:, :1]
+            presence_logits = self.presence_head(self.presence_layer_norm(presence_hidden)).squeeze(-1)
+            presence_logits = presence_logits.clamp(
+                min=-self.clamp_presence_logit_max_val, max=self.clamp_presence_logit_max_val
+            )
+            intermediate_presence_logits.append(presence_logits)
 
         # Stack outputs from all layers
         intermediate_outputs = torch.stack(intermediate_outputs)
         intermediate_boxes = torch.stack(intermediate_boxes[:-1])
-        intermediate_presence_logits = (
-            torch.stack(intermediate_presence_logits) if intermediate_presence_logits else None
-        )
+        intermediate_presence_logits = torch.stack(intermediate_presence_logits)
 
         return Sam3DETRDecoderOutput(
             intermediate_hidden_states=intermediate_outputs,
@@ -1990,6 +2008,8 @@ class Sam3MaskDecoder(Sam3PreTrainedModel):
         self.prompt_cross_attn_norm = nn.LayerNorm(hidden_size)
         self.prompt_cross_attn_dropout = nn.Dropout(config.dropout)
 
+        self.post_init()
+
     @check_model_inputs
     def forward(
         self,

transformers/models/sam3_tracker/modeling_sam3_tracker.py

@@ -107,7 +107,12 @@ class Sam3TrackerFeedForward(nn.Module):
         return hidden_states
 
 
-@auto_docstring
+@auto_docstring(
+    custom_intro="""
+    Segment Anything Model 3 (SAM 3) for generating segmentation masks, given an input image and
+    input points and labels, boxes, or masks.
+    """
+)
 class Sam3TrackerPreTrainedModel(PreTrainedModel):
     config_class = Sam3TrackerConfig
     base_model_prefix = "sam3_tracker"
@@ -123,6 +128,8 @@ class Sam3TrackerPreTrainedModel(PreTrainedModel):
         if isinstance(module, Sam3TrackerModel):
             if module.no_memory_embedding is not None:
                 init.zeros_(module.no_memory_embedding)
+        elif isinstance(module, Sam3TrackerPositionalEmbedding):
+            init.normal_(module.positional_embedding, std=module.scale)
 
 
 class Sam3TrackerPositionalEmbedding(nn.Module):

transformers/models/sam3_tracker/modular_sam3_tracker.py

@@ -136,7 +136,12 @@ class Sam3TrackerFeedForward(Sam2FeedForward):
     pass
 
 
-@auto_docstring
+@auto_docstring(
+    custom_intro="""
+    Segment Anything Model 3 (SAM 3) for generating segmentation masks, given an input image and
+    input points and labels, boxes, or masks.
+    """
+)
 class Sam3TrackerPreTrainedModel(Sam2PreTrainedModel):
     @torch.no_grad()
     def _init_weights(self, module):
@@ -144,6 +149,8 @@ class Sam3TrackerPreTrainedModel(Sam2PreTrainedModel):
         if isinstance(module, Sam3TrackerModel):
             if module.no_memory_embedding is not None:
                 init.zeros_(module.no_memory_embedding)
+        elif isinstance(module, Sam3TrackerPositionalEmbedding):
+            init.normal_(module.positional_embedding, std=module.scale)
 
 
 class Sam3TrackerPositionalEmbedding(Sam2PositionalEmbedding):

transformers/models/sam3_tracker_video/configuration_sam3_tracker_video.py

@@ -397,5 +397,30 @@ class Sam3TrackerVideoConfig(PreTrainedConfig):
 
         super().__init__(**kwargs)
 
+    @property
+    def image_size(self):
+        """Image size for the tracker video model."""
+        return self.vision_config.image_size
+
+    @image_size.setter
+    def image_size(self, value):
+        """Set the image size and propagate to sub-configs. Calculates feature sizes based on patch_size."""
+        self.prompt_encoder_config.image_size = value
+        self.vision_config.image_size = value
+
+        patch_size = self.vision_config.backbone_config.patch_size
+        self.vision_config.backbone_feature_sizes = [
+            [4 * value // patch_size, 4 * value // patch_size],
+            [2 * value // patch_size, 2 * value // patch_size],
+            [value // patch_size, value // patch_size],
+        ]
+        self.memory_attention_rope_feat_sizes = [
+            value // patch_size,
+            value // patch_size,
+        ]
+
+        # keep the image_size in the __dict__ to save the value in the config file (backward compatibility)
+        self.__dict__["image_size"] = value
+
 
 __all__ = ["Sam3TrackerVideoMaskDecoderConfig", "Sam3TrackerVideoPromptEncoderConfig", "Sam3TrackerVideoConfig"]

transformers/models/sam3_tracker_video/modeling_sam3_tracker_video.py

@@ -213,7 +213,7 @@ class Sam3TrackerVideoInferenceSession:
         device_inputs = {}
         for key, value in inputs.items():
             if isinstance(value, torch.Tensor):
-                device_inputs[key] = value.to(self.inference_device, non_blocking=True)
+                device_inputs[key] = value.to(self.inference_device, non_blocking=False)
             else:
                 device_inputs[key] = value
         self.point_inputs_per_obj[obj_idx][frame_idx] = device_inputs
@@ -692,6 +692,12 @@ class Sam3TrackerVideoPreTrainedModel(PreTrainedModel):
         if isinstance(module, Sam3TrackerVideoMemoryFuserCXBlock):
             if module.scale is not None:
                 init.zeros_(module.scale)
+        elif isinstance(module, Sam3TrackerVideoVisionRotaryEmbedding):
+            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, Sam3TrackerVideoPositionalEmbedding):
+            init.normal_(module.positional_embedding, std=module.scale)
 
 
 class Sam3TrackerVideoVisionRotaryEmbedding(nn.Module):
@@ -702,24 +708,17 @@ class Sam3TrackerVideoVisionRotaryEmbedding(nn.Module):
 
     def __init__(self, config: Sam3TrackerVideoConfig):
         super().__init__()
-        dim = config.memory_attention_hidden_size // (
+        self.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 dim % 4 != 0:
+        if self.dim % 4 != 0:
             raise ValueError("Dimension must be divisible by 4 for axial RoPE")
-        end_x, end_y = config.memory_attention_rope_feat_sizes
-        freqs = 1.0 / (config.memory_attention_rope_theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
+        self.end_x, self.end_y = config.memory_attention_rope_feat_sizes
+        self.memory_attention_rope_theta = config.memory_attention_rope_theta
 
-        # 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)
 
@@ -728,6 +727,20 @@ class Sam3TrackerVideoVisionRotaryEmbedding(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 rotate_pairwise(x):
     """
@@ -1567,8 +1580,6 @@ class Sam3TrackerVideoModel(Sam3TrackerVideoPreTrainedModel):
     input_modalities = ("video", "text")
     _can_record_outputs = {"mask_decoder_attentions": OutputRecorder(Sam3TrackerVideoTwoWayAttentionBlock, index=2)}
     _keys_to_ignore_on_load_unexpected = [r"^detector_model."]
-    _tied_weights_keys = {}
-    _keys_to_ignore_on_load_missing = []
     _checkpoint_conversion_mapping = {
         r"tracker_model.(.+)": r"\1",  # the regex allows to remove the prefix, and add it back in revert mode
         "detector_model.vision_encoder.backbone.": "vision_encoder.backbone.",
@@ -1719,6 +1730,7 @@ class Sam3TrackerVideoModel(Sam3TrackerVideoPreTrainedModel):
         frame: Optional[torch.Tensor] = None,
         reverse: bool = False,
         run_mem_encoder: bool = True,
+        **kwargs,
     ) -> Sam3TrackerVideoSegmentationOutput:
         r"""
         inference_session (`Sam3TrackerVideoInferenceSession`):

transformers/models/sam3_tracker_video/modular_sam3_tracker_video.py

@@ -353,6 +353,31 @@ class Sam3TrackerVideoConfig(PreTrainedConfig):
 
         super().__init__(**kwargs)
 
+    @property
+    def image_size(self):
+        """Image size for the tracker video model."""
+        return self.vision_config.image_size
+
+    @image_size.setter
+    def image_size(self, value):
+        """Set the image size and propagate to sub-configs. Calculates feature sizes based on patch_size."""
+        self.prompt_encoder_config.image_size = value
+        self.vision_config.image_size = value
+
+        patch_size = self.vision_config.backbone_config.patch_size
+        self.vision_config.backbone_feature_sizes = [
+            [4 * value // patch_size, 4 * value // patch_size],
+            [2 * value // patch_size, 2 * value // patch_size],
+            [value // patch_size, value // patch_size],
+        ]
+        self.memory_attention_rope_feat_sizes = [
+            value // patch_size,
+            value // patch_size,
+        ]
+
+        # keep the image_size in the __dict__ to save the value in the config file (backward compatibility)
+        self.__dict__["image_size"] = value
+
 
 class Sam3TrackerVideoInferenceCache(Sam2VideoInferenceCache):
     pass
@@ -461,8 +486,6 @@ class Sam3TrackerVideoModel(Sam2VideoModel):
         "tracker_neck.": "vision_encoder.neck.",
     }
     _keys_to_ignore_on_load_unexpected = [r"^detector_model."]
-    _tied_weights_keys = {}
-    _keys_to_ignore_on_load_missing = []
 
     def __init__(self, config: Sam3TrackerVideoConfig, remove_vision_encoder: bool = False):
         r"""

transformers/models/sam3_video/configuration_sam3_video.py

@@ -96,6 +96,9 @@ class Sam3VideoConfig(PreTrainedConfig):
     >>> # Initializing a SAM3 Video configuration with default detector and tracker
     >>> configuration = Sam3VideoConfig()
 
+    >>> # Changing image size for custom resolution inference (automatically propagates to all nested configs)
+    >>> configuration.image_size = 560
+
     >>> # Initializing a model from the configuration
     >>> model = Sam3VideoModel(configuration)
 
@@ -225,5 +228,16 @@ class Sam3VideoConfig(PreTrainedConfig):
         self.high_conf_thresh = high_conf_thresh
         self.high_iou_thresh = high_iou_thresh
 
+    @property
+    def image_size(self):
+        """Image size for the video model."""
+        return self.detector_config.image_size
+
+    @image_size.setter
+    def image_size(self, value):
+        """Recursively propagate the image size to detector and tracker configs."""
+        self.detector_config.image_size = value
+        self.tracker_config.image_size = value
+
 
 __all__ = ["Sam3VideoConfig"]

transformers/models/sam3_video/modeling_sam3_video.py

@@ -33,7 +33,7 @@ from .configuration_sam3_video import Sam3VideoConfig
 
 
 if is_kernels_available():
-    from kernels import get_kernel
+    from ...integrations.hub_kernels import get_kernel
 
 logger = logging.get_logger(__name__)
 
@@ -505,8 +505,6 @@ class Sam3VideoPreTrainedModel(PreTrainedModel):
 
 @auto_docstring
 class Sam3VideoModel(Sam3VideoPreTrainedModel):
-    all_tied_weights_keys = {}
-
     def __init__(self, config: Sam3VideoConfig):
         super().__init__(config)
         self.config = config
@@ -542,6 +540,8 @@ class Sam3VideoModel(Sam3VideoPreTrainedModel):
 
         self.tracker_neck = Sam3VisionNeck(config.detector_config.vision_config)
 
+        self.post_init()
+
     def get_vision_features_for_tracker(self, vision_embeds: torch.Tensor):
         hidden_states = vision_embeds.last_hidden_state
         batch_size = hidden_states.shape[0]
@@ -1697,6 +1697,7 @@ class Sam3VideoModel(Sam3VideoPreTrainedModel):
         frame_idx: Optional[int] = None,
         frame: Optional[torch.Tensor] = None,
         reverse: bool = False,
+        **kwargs,
     ):
         r"""
         inference_session (`Sam3VideoInferenceSession`):