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

transformers/models/sam3/modeling_sam3.py

@@ -1,3 +1,4 @@
+# coding=utf-8
 # Copyright 2025 The Meta AI Authors and The HuggingFace Team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,9 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
+import collections.abc
 import math
-from collections.abc import Callable, Iterable
+from collections.abc import Callable
 from dataclasses import dataclass
+from typing import Optional, Union
 
 import numpy as np
 import torch
@@ -31,14 +35,13 @@ from ...masking_utils import create_bidirectional_mask
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_outputs import (
     BaseModelOutput,
-    BaseModelOutputWithPooling,
     ModelOutput,
 )
 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, can_return_tuple, logging
-from ...utils.generic import TransformersKwargs, check_model_inputs, is_flash_attention_requested
+from ...utils import auto_docstring
+from ...utils.generic import TransformersKwargs, check_model_inputs
 from ..auto import AutoModel
 from .configuration_sam3 import (
     Sam3Config,
@@ -51,21 +54,25 @@ from .configuration_sam3 import (
 )
 
 
-logger = logging.get_logger(__name__)
-
-
 @dataclass
 @auto_docstring
-class Sam3VisionEncoderOutput(BaseModelOutputWithPooling):
+class Sam3VisionEncoderOutput(ModelOutput):
     r"""
     fpn_hidden_states (`tuple[torch.FloatTensor]`):
         Tuple of multi-level FPN feature maps.
     fpn_position_encoding (`tuple[torch.FloatTensor]`):
         Tuple of position encodings for each FPN level.
+    hidden_states (`tuple[torch.FloatTensor]`, *optional*):
+        Tuple of hidden states from all ViT layers.
+    attentions (`tuple[torch.FloatTensor]`, *optional*):
+        Tuple of attention weights from all ViT layers.
     """
 
+    last_hidden_state: torch.FloatTensor = None
     fpn_hidden_states: tuple[torch.FloatTensor, ...] = None
     fpn_position_encoding: tuple[torch.FloatTensor, ...] = None
+    hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    attentions: Optional[tuple[torch.FloatTensor]] = None
 
 
 @dataclass
@@ -79,7 +86,7 @@ class Sam3GeometryEncoderOutput(ModelOutput):
     """
 
     last_hidden_state: torch.FloatTensor = None
-    attention_mask: torch.BoolTensor | None = None
+    attention_mask: Optional[torch.BoolTensor] = None
 
 
 @dataclass
@@ -101,11 +108,11 @@ class Sam3DETREncoderOutput(ModelOutput):
     """
 
     last_hidden_state: torch.FloatTensor = None
-    pos_embeds_flattened: torch.FloatTensor | None = None
-    text_features: torch.FloatTensor | None = None
-    spatial_shapes: torch.LongTensor | None = None
-    hidden_states: tuple[torch.FloatTensor] | None = None
-    attentions: tuple[torch.FloatTensor] | None = None
+    pos_embeds_flattened: Optional[torch.FloatTensor] = None
+    text_features: Optional[torch.FloatTensor] = None
+    spatial_shapes: Optional[torch.LongTensor] = None
+    hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    attentions: Optional[tuple[torch.FloatTensor]] = None
 
 
 @dataclass
@@ -116,8 +123,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, 1)`):
-        Presence logits from all decoder layers indicating object presence confidence.
+    presence_logits (`torch.FloatTensor` of shape `(num_layers, batch_size)`, *optional*):
+        Presence logits from all decoder layers (None if using instance queries).
     hidden_states (`tuple[torch.FloatTensor]`, *optional*):
         Tuple of hidden states from all decoder layers.
     attentions (`tuple[torch.FloatTensor]`, *optional*):
@@ -126,9 +133,9 @@ class Sam3DETRDecoderOutput(ModelOutput):
 
     intermediate_hidden_states: torch.FloatTensor = None
     reference_boxes: torch.FloatTensor = None
-    presence_logits: torch.FloatTensor = None
-    hidden_states: tuple[torch.FloatTensor] | None = None
-    attentions: tuple[torch.FloatTensor] | None = None
+    presence_logits: Optional[torch.FloatTensor] = None
+    hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    attentions: Optional[tuple[torch.FloatTensor]] = None
 
 
 @dataclass
@@ -144,8 +151,8 @@ class Sam3MaskDecoderOutput(ModelOutput):
     """
 
     pred_masks: torch.FloatTensor = None
-    semantic_seg: torch.FloatTensor | None = None
-    attentions: tuple[torch.FloatTensor] | None = None
+    semantic_seg: Optional[torch.FloatTensor] = None
+    attentions: Optional[tuple[torch.FloatTensor]] = None
 
 
 @dataclass
@@ -185,17 +192,17 @@ class Sam3ImageSegmentationOutput(ModelOutput):
 
     pred_masks: torch.FloatTensor = None
     pred_boxes: torch.FloatTensor = None
-    pred_logits: torch.FloatTensor | None = None
-    presence_logits: torch.FloatTensor | None = None
-    semantic_seg: torch.FloatTensor | None = None
-    decoder_hidden_states: tuple[torch.FloatTensor] | None = None
-    decoder_reference_boxes: torch.FloatTensor | None = None
-    encoder_hidden_states: tuple[torch.FloatTensor] | None = None
-    vision_hidden_states: tuple[torch.FloatTensor] | None = None
-    vision_attentions: tuple[torch.FloatTensor] | None = None
-    detr_encoder_attentions: tuple[torch.FloatTensor] | None = None
-    detr_decoder_attentions: tuple[torch.FloatTensor] | None = None
-    mask_decoder_attentions: tuple[torch.FloatTensor] | None = None
+    pred_logits: Optional[torch.FloatTensor] = None
+    presence_logits: Optional[torch.FloatTensor] = None
+    semantic_seg: Optional[torch.FloatTensor] = None
+    decoder_hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    decoder_reference_boxes: Optional[torch.FloatTensor] = None
+    encoder_hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    vision_hidden_states: Optional[tuple[torch.FloatTensor]] = None
+    vision_attentions: Optional[tuple[torch.FloatTensor]] = None
+    detr_encoder_attentions: Optional[tuple[torch.FloatTensor]] = None
+    detr_decoder_attentions: Optional[tuple[torch.FloatTensor]] = None
+    mask_decoder_attentions: Optional[tuple[torch.FloatTensor]] = None
 
 
 def inverse_sigmoid(x: torch.Tensor, eps: float = 1e-3) -> torch.Tensor:
@@ -268,7 +275,7 @@ def box_cxcywh_to_xyxy(x):
 
 
 class Sam3MLP(nn.Module):
-    def __init__(self, config: Sam3ViTConfig):
+    def __init__(self, config: Union[Sam3ViTConfig]):
         super().__init__()
         self.config = config
         self.activation_fn = ACT2FN[config.hidden_act]
@@ -289,8 +296,8 @@ def eager_attention_forward(
     query: torch.Tensor,
     key: torch.Tensor,
     value: torch.Tensor,
-    attention_mask: torch.Tensor | None,
-    scaling: float | None = None,
+    attention_mask: Optional[torch.Tensor],
+    scaling: Optional[float] = None,
     dropout: float = 0.0,
     **kwargs: Unpack[TransformersKwargs],
 ):
@@ -338,7 +345,7 @@ class Sam3Attention(nn.Module):
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
+        attention_mask: Optional[torch.Tensor] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
@@ -365,19 +372,6 @@ class Sam3Attention(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_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,
@@ -408,10 +402,6 @@ 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)
@@ -541,8 +531,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, Iterable) else (image_size, image_size)
-        patch_size = patch_size if isinstance(patch_size, Iterable) else (patch_size, patch_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)
         num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
         self.image_size = image_size
         self.patch_size = patch_size
@@ -552,7 +542,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.to(self.projection.weight.dtype)).flatten(2).transpose(1, 2)
+        embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2)
         return embeddings
 
 
@@ -771,19 +761,6 @@ 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
@@ -838,7 +815,7 @@ class Sam3SinePositionEmbedding(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:
@@ -908,9 +885,9 @@ class Sam3SinePositionEmbedding(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)
@@ -961,7 +938,6 @@ 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)
 
@@ -1029,9 +1005,9 @@ class Sam3VisionModel(Sam3PreTrainedModel):
     @check_model_inputs
     def forward(
         self,
-        pixel_values: torch.FloatTensor | None = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
         **kwargs: Unpack[TransformersKwargs],
-    ) -> tuple | Sam3VisionEncoderOutput:
+    ) -> Union[tuple, Sam3VisionEncoderOutput]:
         if pixel_values is None:
             raise ValueError("You have to specify pixel_values")
 
@@ -1191,7 +1167,7 @@ class Sam3GeometryEncoder(nn.Module):
         box_mask: torch.Tensor,
         box_labels: torch.Tensor,
         img_feats: tuple[torch.Tensor, ...],
-        img_pos_embeds: tuple[torch.Tensor, ...] | None = None,
+        img_pos_embeds: Optional[tuple[torch.Tensor, ...]] = None,
     ):
         """
         Forward pass for encoding geometric prompts.
@@ -1277,7 +1253,7 @@ class Sam3DetrEncoderLayer(nn.Module):
         vision_feats: Tensor,
         prompt_feats: Tensor,
         vision_pos_encoding: Tensor,
-        prompt_cross_attn_mask: Tensor | None = None,
+        prompt_mask: Tensor,
         **kwargs: Unpack[TransformersKwargs],
     ):
         """
@@ -1287,7 +1263,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_cross_attn_mask: Cross-attention mask for prompt features
+            prompt_mask: Padding mask for prompts [batch_size, text_len] where True=valid, False=padding
 
         Returns:
             Updated vision features [batch_size, vision_len, hidden_size]
@@ -1308,6 +1284,15 @@ 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,
@@ -1346,8 +1331,6 @@ 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],
@@ -1395,11 +1378,11 @@ class Sam3DetrEncoder(Sam3PreTrainedModel):
         self,
         vision_features: list[torch.Tensor],
         text_features: torch.Tensor,
-        vision_pos_embeds: list[torch.Tensor] | None = None,
-        text_mask: torch.Tensor | None = None,
-        spatial_sizes: list[tuple[int, int]] | None = None,
+        vision_pos_embeds: Optional[list[torch.Tensor]] = None,
+        text_mask: Optional[torch.Tensor] = None,
+        spatial_sizes: Optional[list[tuple[int, int]]] = None,
         **kwargs: Unpack[TransformersKwargs],
-    ) -> tuple | Sam3DETREncoderOutput:
+    ):
         """
         Forward pass for the DETR encoder.
 
@@ -1429,22 +1412,13 @@ 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_cross_attn_mask=prompt_cross_attn_mask,
+                prompt_mask=text_mask,
                 **kwargs,
             )
         return Sam3DETREncoderOutput(
@@ -1510,27 +1484,31 @@ class Sam3DetrDecoderLayer(nn.Module):
         text_features: torch.Tensor,
         vision_features: torch.Tensor,
         vision_pos_encoding: torch.Tensor,
-        text_cross_attn_mask: torch.Tensor | None = None,
-        vision_cross_attn_mask: torch.Tensor | None = None,
+        text_mask: Optional[torch.Tensor] = None,
+        vision_cross_attn_mask: Optional[torch.Tensor] = None,
+        presence_token: Optional[torch.Tensor] = None,
         **kwargs: Unpack[TransformersKwargs],
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         """
         Forward pass for decoder layer.
 
         Args:
-            hidden_states: Query features [batch_size, num_queries + 1, hidden_size] (includes presence token at position 0)
+            hidden_states: Query features [batch_size, num_queries, hidden_size]
             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_cross_attn_mask: Text cross-attention mask
-            vision_cross_attn_mask: Vision cross-attention mask, already expanded for presence token
+            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]
 
         Returns:
-            Updated hidden states (including presence token at position 0)
+            Tuple of (updated hidden states, updated presence token)
         """
-        # Prepend zeros to query_pos for presence token
-        query_pos = F.pad(query_pos, (0, 0, 1, 0), mode="constant", value=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)
 
         # Self-attention with query position encoding
         residual = hidden_states
@@ -1549,6 +1527,15 @@ 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,
@@ -1559,6 +1546,20 @@ 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
@@ -1567,7 +1568,7 @@ class Sam3DetrDecoderLayer(nn.Module):
             query=query_with_pos,
             key=key_with_pos,
             value=vision_features,
-            attention_mask=vision_cross_attn_mask,
+            attention_mask=combined_vision_mask,
             **kwargs,
         )
         hidden_states = residual + self.vision_cross_attn_dropout(attn_output)
@@ -1579,7 +1580,13 @@ class Sam3DetrDecoderLayer(nn.Module):
         hidden_states = residual + self.mlp_dropout(hidden_states)
         hidden_states = self.mlp_layer_norm(hidden_states)
 
-        return 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
 
 
 class Sam3DetrDecoder(Sam3PreTrainedModel):
@@ -1627,8 +1634,6 @@ 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
@@ -1691,10 +1696,10 @@ class Sam3DetrDecoder(Sam3PreTrainedModel):
         vision_features: torch.Tensor,
         text_features: torch.Tensor,
         vision_pos_encoding: torch.Tensor,
-        text_mask: torch.Tensor | None = None,
-        spatial_shapes: torch.Tensor | None = None,
+        text_mask: Optional[torch.Tensor] = None,
+        spatial_shapes: Optional[torch.Tensor] = None,
         **kwargs: Unpack[TransformersKwargs],
-    ) -> tuple | Sam3DETRDecoderOutput:
+    ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
         """
         Forward pass for the DETR decoder.
 
@@ -1710,23 +1715,11 @@ class Sam3DetrDecoder(Sam3PreTrainedModel):
         """
         batch_size = vision_features.shape[0]
 
-        query_embeds = self.query_embed.weight.unsqueeze(0).expand(batch_size, -1, -1)
+        hidden_states = 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 = []
@@ -1741,45 +1734,43 @@ 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])
-                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)
+                vision_cross_attn_mask = self._get_rpb_matrix(reference_boxes, spatial_shape)
 
-            hidden_states = layer(
+            hidden_states, presence_token = layer(
                 hidden_states,
                 query_pos=query_pos,
                 text_features=text_features,
                 vision_features=vision_features,
                 vision_pos_encoding=vision_pos_encoding,
-                text_cross_attn_mask=text_cross_attn_mask,
+                text_mask=text_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(query_hidden_states))
+            delta_boxes = self.box_head(self.output_layer_norm(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(query_hidden_states))
+            intermediate_outputs.append(self.output_layer_norm(hidden_states))
             intermediate_boxes.append(new_reference_boxes)
 
             # Process presence token
-            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)
+            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)
 
         # 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)
+        intermediate_presence_logits = (
+            torch.stack(intermediate_presence_logits) if intermediate_presence_logits else None
+        )
 
         return Sam3DETRDecoderOutput(
             intermediate_hidden_states=intermediate_outputs,
@@ -1820,7 +1811,7 @@ class Sam3DotProductScoring(nn.Module):
         self.clamp_logits = True
         self.clamp_max_val = 12.0
 
-    def _pool_text_features(self, text_features: torch.Tensor, text_mask: torch.Tensor | None) -> torch.Tensor:
+    def _pool_text_features(self, text_features: torch.Tensor, text_mask: Optional[torch.Tensor]) -> torch.Tensor:
         """
         Mean pool text features, accounting for padding.
 
@@ -1849,7 +1840,7 @@ class Sam3DotProductScoring(nn.Module):
         self,
         decoder_hidden_states: torch.Tensor,
         text_features: torch.Tensor,
-        text_mask: torch.Tensor | None = None,
+        text_mask: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         """
         Compute classification scores via dot product.
@@ -1999,18 +1990,16 @@ 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,
         decoder_queries: torch.Tensor,
         backbone_features: list[torch.Tensor],
         encoder_hidden_states: torch.Tensor,
-        prompt_features: torch.Tensor | None = None,
-        prompt_mask: torch.Tensor | None = None,
+        prompt_features: Optional[torch.Tensor] = None,
+        prompt_mask: Optional[torch.Tensor] = None,
         **kwargs: Unpack[TransformersKwargs],
-    ) -> tuple | Sam3MaskDecoderOutput:
+    ) -> dict[str, torch.Tensor]:
         """
         Args:
             decoder_queries: Decoder output queries [batch_size, num_queries, hidden_size]
@@ -2140,45 +2129,44 @@ class Sam3Model(Sam3PreTrainedModel):
 
         self.post_init()
 
-    @can_return_tuple
     @auto_docstring
     def get_text_features(
         self,
         input_ids: torch.LongTensor,
-        attention_mask: torch.Tensor | None = None,
+        attention_mask: Optional[torch.Tensor] = None,
         **kwargs: Unpack[TransformersKwargs],
-    ) -> tuple | BaseModelOutputWithPooling:
+    ) -> torch.FloatTensor:
         r"""
+        Returns:
+            text_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+                Text embeddings that can be passed as `text_embeds` to the forward method.
+
         Example:
 
         ```python
         >>> from transformers import Sam3Model, Sam3Processor
         >>> from PIL import Image
-        >>> import httpx
-        >>> from io import BytesIO
+        >>> import requests
 
         >>> model = Sam3Model.from_pretrained("facebook/sam3")
         >>> processor = Sam3Processor.from_pretrained("facebook/sam3")
 
         >>> # Pre-compute text embeddings
         >>> text_inputs = processor(text="cat", return_tensors="pt")
-        >>> text_embeds = model.get_text_features(**text_inputs).pooler_output
+        >>> text_embeds = model.get_text_features(**text_inputs)
 
         >>> # Reuse text embeddings for multiple images
-        >>> url = "http://images.cocodataset.org/val2017/000000077595.jpg"
-        >>> with httpx.stream("GET", url) as response:
-        ...     image = Image.open(BytesIO(response.read()))
+        >>> img_url = "http://images.cocodataset.org/val2017/000000077595.jpg"
+        >>> image = Image.open(requests.get(img_url, stream=True).raw)
         >>> img_inputs = processor(images=image, return_tensors="pt")
         >>> outputs = model(pixel_values=img_inputs.pixel_values, text_embeds=text_embeds)
         ```
         """
-        text_outputs = self.text_encoder(
-            input_ids=input_ids, attention_mask=attention_mask, return_dict=True, **kwargs
-        )
-        last_hidden_state = text_outputs.last_hidden_state
-        text_outputs.pooler_output = self.text_projection(last_hidden_state)
-
-        return text_outputs
+        text_features = self.text_encoder(
+            input_ids=input_ids, attention_mask=attention_mask, **kwargs
+        ).last_hidden_state
+        text_features = self.text_projection(text_features)
+        return text_features
 
     @auto_docstring
     def get_vision_features(
@@ -2187,21 +2175,23 @@ class Sam3Model(Sam3PreTrainedModel):
         **kwargs: Unpack[TransformersKwargs],
     ) -> Sam3VisionEncoderOutput:
         r"""
+        Returns:
+            vision_embeds (`Sam3VisionEncoderOutput`):
+                Vision embeddings that can be passed as `vision_embeds` to the forward method.
+
         Example:
 
         ```python
         >>> from transformers import Sam3Model, Sam3Processor
         >>> from PIL import Image
-        >>> import httpx
-        >>> from io import BytesIO
+        >>> import requests
 
         >>> model = Sam3Model.from_pretrained("facebook/sam3")
         >>> processor = Sam3Processor.from_pretrained("facebook/sam3")
 
         >>> # Pre-compute vision embeddings
-        >>> url = "http://images.cocodataset.org/val2017/000000077595.jpg"
-        >>> with httpx.stream("GET", url) as response:
-        ...     image = Image.open(BytesIO(response.read()))
+        >>> img_url = "http://images.cocodataset.org/val2017/000000077595.jpg"
+        >>> image = Image.open(requests.get(img_url, stream=True).raw)
         >>> img_inputs = processor(images=image, return_tensors="pt")
         >>> vision_embeds = model.get_vision_features(pixel_values=img_inputs.pixel_values)
 
@@ -2217,13 +2207,13 @@ class Sam3Model(Sam3PreTrainedModel):
     @auto_docstring
     def forward(
         self,
-        pixel_values: torch.FloatTensor | None = None,
-        vision_embeds: Sam3VisionEncoderOutput | None = None,
-        input_ids: torch.LongTensor | None = None,
-        attention_mask: torch.Tensor | None = None,
-        text_embeds: torch.FloatTensor | None = None,
-        input_boxes: torch.FloatTensor | None = None,
-        input_boxes_labels: torch.LongTensor | None = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        vision_embeds: Optional[Sam3VisionEncoderOutput] = None,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        text_embeds: Optional[torch.FloatTensor] = None,
+        input_boxes: Optional[torch.FloatTensor] = None,
+        input_boxes_labels: Optional[torch.LongTensor] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> Sam3ImageSegmentationOutput:
         r"""
@@ -2242,18 +2232,16 @@ class Sam3Model(Sam3PreTrainedModel):
 
         ```python
         >>> from PIL import Image
-        >>> import httpx
-        >>> from io import BytesIO
+        >>> import requests
         >>> from transformers import AutoModel, AutoProcessor
 
         >>> model = AutoModel.from_pretrained("facebook/sam3")
         >>> processor = AutoProcessor.from_pretrained("facebook/sam3")
 
-        >>> url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/sam-car.png"
-        >>> with httpx.stream("GET", url) as response:
-        ...     image = Image.open(BytesIO(response.read())).convert("RGB")
+        >>> img_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/sam-car.png"
+        >>> raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
         >>> text = "car"
-        >>> inputs = processor(images=image, text=text, return_tensors="pt")
+        >>> inputs = processor(images=raw_image, text=text, return_tensors="pt")
 
         >>> # Get segmentation output
         >>> outputs = model(**inputs)
@@ -2283,9 +2271,7 @@ class Sam3Model(Sam3PreTrainedModel):
         fpn_position_encoding = vision_outputs.fpn_position_encoding[:-1]
 
         if text_embeds is None:
-            text_features = self.get_text_features(
-                input_ids=input_ids, attention_mask=attention_mask, return_dict=True
-            ).pooler_output
+            text_features = self.get_text_features(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
         else:
             text_features = text_embeds