import math

from collections.abc import Callable, Iterable

from dataclasses import dataclass

import numpy as np

import torch

import torch.nn as nn

import torch.nn.functional as F

from torch import Tensor

from ...utils import is_torchvision_available

if is_torchvision_available():

import torchvision

from transformers import CLIPTextModelWithProjection

from ... import initialization as init

from ...activations import ACT2FN

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,

is_flash_attention_requested,

merge_with_config_defaults,

)

from ...utils.import_utils import requires

from ...utils.output_capturing import capture_outputs

from ..auto import AutoModel

from .configuration_sam3 import (

Sam3Config,

Sam3DETRDecoderConfig,

Sam3DETREncoderConfig,

Sam3GeometryEncoderConfig,

Sam3MaskDecoderConfig,

Sam3VisionConfig,

Sam3ViTConfig,

)

logger = logging.get_logger(__name__)

class Sam3VisionEncoderOutput(BaseModelOutputWithPooling):

r"""

fpn_hidden_states: tuple[torch.FloatTensor, ...] = None

fpn_position_encoding: tuple[torch.FloatTensor, ...] = None

class Sam3GeometryEncoderOutput(ModelOutput):

r"""

last_hidden_state: torch.FloatTensor = None

attention_mask: torch.BoolTensor | None = None

class Sam3DETREncoderOutput(ModelOutput):

r"""

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

class Sam3DETRDecoderOutput(ModelOutput):

r"""

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

class Sam3MaskDecoderOutput(ModelOutput):

r"""

pred_masks: torch.FloatTensor = None

semantic_seg: torch.FloatTensor | None = None

attentions: tuple[torch.FloatTensor] | None = None

class Sam3ImageSegmentationOutput(ModelOutput):

r"""

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

def inverse_sigmoid(x: torch.Tensor, eps: float = 1e-3) -> torch.Tensor:

"""The inverse function for sigmoid activation function."""

x = x.clamp(min=0, max=1)

x1 = x.clamp(min=eps)

x2 = (1 - x).clamp(min=eps)

return torch.log(x1 / x2)

def concat_padded_sequences(seq1, mask1, seq2, mask2, return_index: bool = False):

"""

batch_size, seq1_length, hidden_size = seq1.shape

batch_size2, seq2_length, hidden_size2 = seq2.shape

assert batch_size == batch_size2 == mask1.size(0) == mask2.size(0)

assert hidden_size == hidden_size2

assert seq1_length == mask1.size(1)

assert seq2_length == mask2.size(1)

actual_seq1_lengths = mask1.sum(dim=-1)

actual_seq2_lengths = mask2.sum(dim=-1)

final_lengths = actual_seq1_lengths + actual_seq2_lengths

max_length = seq1_length + seq2_length

concatenated_mask = (

torch.arange(max_length, device=seq2.device)[None].repeat(batch_size, 1) < final_lengths[:, None]

)

concatenated_sequence = torch.zeros((batch_size, max_length, hidden_size), device=seq2.device, dtype=seq2.dtype)

concatenated_sequence[:, :seq1_length, :] = seq1

# Shift seq2 elements to start at the end of valid seq1

index = torch.arange(seq2_length, device=seq2.device)[None].repeat(batch_size, 1)

index = index + actual_seq1_lengths[:, None]

# Scatter seq2 into the right positions

concatenated_sequence = concatenated_sequence.scatter(1, index[:, :, None].expand(-1, -1, hidden_size), seq2)

if return_index:

return concatenated_sequence, concatenated_mask, index

return concatenated_sequence, concatenated_mask

def box_cxcywh_to_xyxy(x):

"""Convert boxes from (cx, cy, w, h) format to (x1, y1, x2, y2) format."""

x_c, y_c, w, h = x.unbind(-1)

b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]

return torch.stack(b, dim=-1)

class Sam3MLP(nn.Module):

def __init__(self, config):

super().__init__()

self.config = config

self.activation_fn = ACT2FN[config.hidden_act]

self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)

self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)

self.dropout = nn.Dropout(config.hidden_dropout)

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

hidden_states = self.fc1(hidden_states)

hidden_states = self.dropout(hidden_states)

hidden_states = self.activation_fn(hidden_states)

hidden_states = self.fc2(hidden_states)

return hidden_states

def eager_attention_forward(

module: nn.Module,

query: torch.Tensor,

key: torch.Tensor,

value: torch.Tensor,

attention_mask: torch.Tensor | None,

scaling: float | None = None,

dropout: float = 0.0,

**kwargs: Unpack[TransformersKwargs],

):

if scaling is None:

scaling = query.size(-1) ** -0.5

# Take the dot product between "query" and "key" to get the raw attention scores.

attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling

if attention_mask is not None:

attn_weights = attn_weights + attention_mask

attn_weights = nn.functional.softmax(attn_weights, dim=-1)

attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)

attn_output = torch.matmul(attn_weights, value)

attn_output = attn_output.transpose(1, 2).contiguous()

return attn_output, attn_weights

class Sam3Attention(nn.Module):

"""

def __init__(self, config):

super().__init__()

self.config = config

self.hidden_size = config.hidden_size

self.num_attention_heads = config.num_attention_heads

self.head_dim = self.hidden_size // config.num_attention_heads

self.scaling = self.head_dim**-0.5

self.is_causal = False

self.q_proj = nn.Linear(self.hidden_size, self.hidden_size)

self.k_proj = nn.Linear(self.hidden_size, self.hidden_size)

self.v_proj = nn.Linear(self.hidden_size, self.hidden_size)

self.o_proj = nn.Linear(self.hidden_size, self.hidden_size)

def forward(

self,

query: torch.Tensor,

key: torch.Tensor,

value: torch.Tensor,

attention_mask: torch.Tensor | None = None,

**kwargs: Unpack[TransformersKwargs],

) -> tuple[torch.Tensor, torch.Tensor]:

"""

batch_size = query.shape[0]

query_len = query.shape[1]

key_len = key.shape[1]

query = self.q_proj(query).view(batch_size, query_len, self.num_attention_heads, self.head_dim).transpose(1, 2)

key = self.k_proj(key).view(batch_size, key_len, self.num_attention_heads, self.head_dim).transpose(1, 2)

value = self.v_proj(value).view(batch_size, key_len, self.num_attention_heads, self.head_dim).transpose(1, 2)

attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(

self.config._attn_implementation, eager_attention_forward

)

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,

key,

value,

attention_mask=attention_mask,

dropout=0.0,

scaling=self.scaling,

is_causal=self.is_causal,

**kwargs,

)

attn_output = attn_output.reshape(batch_size, query_len, self.num_attention_heads * self.head_dim).contiguous()

attn_output = self.o_proj(attn_output)

return attn_output, attn_weights

class Sam3ViTRotaryEmbedding(nn.Module):

"""

def __init__(self, config: Sam3ViTConfig, end_x: int, end_y: int, scale: float = 1.0):

super().__init__()

dim = config.hidden_size // config.num_attention_heads

# 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)

x_positions = (flattened_indices % end_x) * scale

y_positions = torch.div(flattened_indices, end_x, rounding_mode="floor") * 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)

# directly register the cos and sin embeddings as we have a fixed feature shape

self.rope_embeddings_cos = nn.Buffer(inv_freq.cos(), persistent=False)

self.rope_embeddings_sin = nn.Buffer(inv_freq.sin(), persistent=False)

@torch.no_grad()

def forward(self) -> tuple[torch.Tensor, torch.Tensor]:

# As the feature map size is fixed for each stage, we can just return the pre-computed embeddings.

return self.rope_embeddings_cos, self.rope_embeddings_sin

def rotate_pairwise(x):

"""

x = x.view(*x.shape[:-1], -1, 2)

x1, x2 = x.unbind(dim=-1)

x = torch.stack((-x2, x1), dim=-1)

return x.flatten(start_dim=-2)

def apply_rotary_pos_emb_2d(

q: torch.Tensor,

k: torch.Tensor,

cos: torch.Tensor,

sin: torch.Tensor,

) -> tuple[torch.Tensor, torch.Tensor]:

"""

q_embed = q.float()

q_embed = (q_embed * cos) + (rotate_pairwise(q_embed) * sin)

k_embed = k.float()

k_embed = (k_embed * cos) + (rotate_pairwise(k_embed) * sin)

return q_embed.type_as(q), k_embed.type_as(k)

class Sam3ViTRoPEAttention(nn.Module):

"""Self-attention with rotary position encoding."""

def __init__(self, config: Sam3ViTConfig):

super().__init__()

self.config = config

self.hidden_size = config.hidden_size

self.num_attention_heads = config.num_attention_heads

self.head_dim = self.hidden_size // config.num_attention_heads

self.scaling = self.head_dim**-0.5

self.attention_dropout = config.attention_dropout

self.is_causal = False

self.q_proj = nn.Linear(self.hidden_size, self.hidden_size)

self.k_proj = nn.Linear(self.hidden_size, self.hidden_size)

self.v_proj = nn.Linear(self.hidden_size, self.hidden_size)

self.o_proj = nn.Linear(self.hidden_size, self.hidden_size)

def forward(

self,

hidden_states: torch.Tensor,

position_embeddings: tuple[torch.Tensor, torch.Tensor],

**kwargs: Unpack[TransformersKwargs],

) -> Tensor:

batch_size, height, width, _ = hidden_states.shape

seq_len = height * width

new_shape = (batch_size, seq_len, self.num_attention_heads, self.head_dim)

query = self.q_proj(hidden_states).view(*new_shape).transpose(1, 2)

key = self.k_proj(hidden_states).view(*new_shape).transpose(1, 2)

value = self.v_proj(hidden_states).view(*new_shape).transpose(1, 2)

cos, sin = position_embeddings

query, key = apply_rotary_pos_emb_2d(query, key, cos=cos, sin=sin)

attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(

self.config._attn_implementation, eager_attention_forward

)

attn_output, attn_weights = attention_interface(

self,

query,

key,

value,

attention_mask=None,

dropout=0.0 if not self.training else self.attention_dropout,

scaling=self.scaling,

is_causal=self.is_causal,

**kwargs,

)

attn_output = attn_output.reshape(batch_size, height, width, -1).contiguous()

attn_output = self.o_proj(attn_output)

return attn_output, attn_weights

class Sam3ViTPatchEmbeddings(nn.Module):

"""

def __init__(self, config: Sam3ViTConfig):

super().__init__()

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)

num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

self.image_size = image_size

self.patch_size = patch_size

self.num_channels = num_channels

self.num_patches = num_patches

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)

return embeddings

class Sam3ViTEmbeddings(nn.Module):

"""

def __init__(self, config: Sam3ViTConfig):

super().__init__()

self.patch_embeddings = Sam3ViTPatchEmbeddings(config)

num_patches = self.patch_embeddings.num_patches

self.position_embeddings = nn.Parameter(

torch.randn(1, num_patches, config.hidden_size)

) # !Remove cls token in convert weights!

self.dropout = nn.Dropout(config.hidden_dropout)

self.patch_size = config.patch_size

def _tile_position_embeddings(

self,

position_embeddings: torch.Tensor,

height: int,

width: int,

) -> torch.Tensor:

"""

pretrain_size = int(position_embeddings.shape[1] ** 0.5)

# Skip tiling if sizes match (but always tile during tracing for consistent graph)

if not torch.jit.is_tracing() and pretrain_size == height and pretrain_size == width:

return position_embeddings.reshape(1, height * width, -1)

# Tile position embeddings to match target spatial dimensions

hidden_size = position_embeddings.shape[-1]

pos_embed = position_embeddings.reshape(1, pretrain_size, pretrain_size, hidden_size).permute(0, 3, 1, 2)

repeat_h = height // pretrain_size + 1

repeat_w = width // pretrain_size + 1

pos_embed = pos_embed.tile([1, 1, repeat_h, repeat_w])[:, :, :height, :width]

return pos_embed.permute(0, 2, 3, 1).reshape(1, height * width, hidden_size)

def forward(

self,

pixel_values: torch.Tensor,

interpolate_pos_encoding: bool = False,

) -> torch.Tensor:

height, width = pixel_values.shape[-2:]

embeddings = self.patch_embeddings(pixel_values)

# Calculate spatial dimensions in patches

height_patches = height // self.patch_size

width_patches = width // self.patch_size

position_embeddings = self._tile_position_embeddings(

self.position_embeddings,

height_patches,

width_patches,

)

embeddings = embeddings + position_embeddings

embeddings = self.dropout(embeddings)

return embeddings

def window_partition(hidden_state, window_size):

"""

batch_size, height, width, num_channels = hidden_state.shape

pad_height = (window_size - height % window_size) % window_size

pad_width = (window_size - width % window_size) % window_size

# Noop in case pad_width == 0 and pad_height == 0.

hidden_state = nn.functional.pad(hidden_state, (0, 0, 0, pad_width, 0, pad_height))

padded_height, padded_width = height + pad_height, width + pad_width

hidden_state = hidden_state.view(

batch_size, padded_height // window_size, window_size, padded_width // window_size, window_size, num_channels

)

windows = hidden_state.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, num_channels)

return windows, (padded_height, padded_width)

def window_unpartition(windows, window_size, pad_height_width, height_width):

"""

padded_height, padded_width = pad_height_width

height, width = height_width

batch_size = windows.shape[0] // (padded_height * padded_width // window_size // window_size)

hidden_state = windows.view(

batch_size, padded_height // window_size, padded_width // window_size, window_size, window_size, -1

)

hidden_state = hidden_state.permute(0, 1, 3, 2, 4, 5).contiguous()

hidden_state = hidden_state.view(batch_size, padded_height, padded_width, -1)

# We always have height <= padded_height and width <= padded_width

hidden_state = hidden_state[:, :height, :width, :].contiguous()

return hidden_state

class Sam3ViTLayerScale(nn.Module):

def __init__(self, config) -> None:

super().__init__()

self.lambda1 = nn.Parameter(config.layer_scale_init_value * torch.ones(config.hidden_size))

def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:

return hidden_state * self.lambda1

class Sam3ViTLayer(GradientCheckpointingLayer):

"""Vision Transformer layer with rotary position embeddings and optional windowed attention."""

def __init__(self, config: Sam3ViTConfig, window_size: int = 0) -> None:

super().__init__()

hidden_size = config.hidden_size

image_size = config.image_size

image_size = image_size if isinstance(image_size, (list, tuple)) else (image_size, image_size)

patch_size = config.patch_size

patch_size = patch_size if isinstance(patch_size, (list, tuple)) else (patch_size, patch_size)

input_size = (image_size[0] // patch_size[0], image_size[1] // patch_size[1])

self.layer_norm1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps)

rotary_input_size = input_size if window_size == 0 else (window_size, window_size)

rotary_scale = config.window_size / rotary_input_size[0]

self.rotary_emb = Sam3ViTRotaryEmbedding(

config, end_x=rotary_input_size[0], end_y=rotary_input_size[1], scale=rotary_scale

)

self.attention = Sam3ViTRoPEAttention(config)

self.layer_norm2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps)

self.mlp = Sam3MLP(config)

self.dropout = nn.Dropout(config.hidden_dropout)

self.window_size = window_size

def forward(

self,

hidden_states: torch.Tensor,

**kwargs: Unpack[TransformersKwargs],

) -> torch.Tensor:

residual = hidden_states

hidden_states = self.layer_norm1(hidden_states)

if self.window_size > 0:

height, width = hidden_states.shape[1], hidden_states.shape[2]

# Partition into non-overlapping windows for efficient attention

hidden_states, pad_height_width = window_partition(hidden_states, self.window_size)

position_embeddings = self.rotary_emb()

hidden_states, _ = self.attention(hidden_states, position_embeddings, **kwargs)

if self.window_size > 0:

# Reverse window partition to restore original spatial layout

hidden_states = window_unpartition(hidden_states, self.window_size, pad_height_width, (height, width))

hidden_states = residual + hidden_states

residual = hidden_states

hidden_states = self.layer_norm2(hidden_states)

hidden_states = self.mlp(hidden_states)

hidden_states = residual + self.dropout(hidden_states)

return hidden_states

class Sam3PreTrainedModel(PreTrainedModel):

config_class = Sam3Config

base_model_prefix = "sam3"

main_input_name = "pixel_values"

input_modalities = ["image", "text"]

_supports_sdpa = True

_supports_flash_attn = True

_supports_flex_attn = True

_supports_attention_backend = True

def _init_weights(self, module):

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())

class Sam3ViTModel(Sam3PreTrainedModel):

config: Sam3ViTConfig

_can_record_outputs = {

"hidden_states": Sam3ViTLayer,

"attentions": Sam3ViTRoPEAttention,

}

def __init__(self, config: Sam3ViTConfig):

super().__init__(config)

self.config = config

self.embeddings = Sam3ViTEmbeddings(config)

self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

self.layers = nn.ModuleList(

[

Sam3ViTLayer(config, window_size=config.window_size if i not in config.global_attn_indexes else 0)

for i in range(config.num_hidden_layers)

]

)

self.post_init()

def get_input_embeddings(self) -> Sam3ViTPatchEmbeddings:

return self.embeddings.patch_embeddings

@merge_with_config_defaults

@capture_outputs(tie_last_hidden_states=False)

@auto_docstring

def forward(

self,

pixel_values: torch.Tensor,

**kwargs: Unpack[TransformersKwargs],

) -> BaseModelOutput:

hidden_states = self.embeddings(pixel_values) # [batch_size, seq_len, hidden_size]

batch_size = hidden_states.shape[0]

height = pixel_values.shape[-2] // self.config.patch_size

width = pixel_values.shape[-1] // self.config.patch_size

hidden_size = hidden_states.shape[-1]

# Reshape to spatial format for windowed attention: [batch_size, height, width, hidden_size]

hidden_states = hidden_states.view(batch_size, height, width, hidden_size)

hidden_states = self.layer_norm(hidden_states)

for layer in self.layers:

hidden_states = layer(hidden_states, **kwargs)

# Reshape back to sequence format: [batch_size, height*width, hidden_size]

hidden_states = hidden_states.view(batch_size, height * width, hidden_size)

return BaseModelOutput(last_hidden_state=hidden_states)

class Sam3SinePositionEmbedding(nn.Module):

"""

def __init__(

self,

num_position_features: int = 64,

temperature: int = 10000,

normalize: bool = False,

scale: float | None = None,

):

super().__init__()

if scale is not None and normalize is False:

raise ValueError("normalize should be True if scale is passed")

self.num_position_features = num_position_features

self.temperature = temperature

self.normalize = normalize

self.scale = 2 * math.pi if scale is None else scale

def encode_1d_positions(self, x: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

"""

x_embed = x * self.scale

y_embed = y * self.scale

dim_t = torch.arange(self.num_position_features, dtype=torch.int64, device=x.device).to(x.dtype)

dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_position_features)

pos_x = x_embed[:, None] / dim_t

pos_y = y_embed[:, None] / dim_t

pos_x = torch.stack((pos_x[:, 0::2].sin(), pos_x[:, 1::2].cos()), dim=2).flatten(1)

pos_y = torch.stack((pos_y[:, 0::2].sin(), pos_y[:, 1::2].cos()), dim=2).flatten(1)

return pos_x, pos_y

def encode_boxes(self, boxes: torch.Tensor) -> torch.Tensor:

"""

assert boxes.size(-1) == 4, f"Expected 4D box coordinates (x, y, w, h), got shape {boxes.shape}"

dim_t = torch.arange(self.num_position_features, dtype=torch.int64, device=boxes.device).to(boxes.dtype)

dim_t = self.temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / self.num_position_features)

x_embed = boxes[:, :, 0] * self.scale

y_embed = boxes[:, :, 1] * self.scale

w_embed = boxes[:, :, 2] * self.scale

h_embed = boxes[:, :, 3] * self.scale

pos_x = x_embed[:, :, None] / dim_t

pos_y = y_embed[:, :, None] / dim_t

pos_w = w_embed[:, :, None] / dim_t

pos_h = h_embed[:, :, None] / dim_t

pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2)

pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2)

pos_w = torch.stack((pos_w[:, :, 0::2].sin(), pos_w[:, :, 1::2].cos()), dim=3).flatten(2)

pos_h = torch.stack((pos_h[:, :, 0::2].sin(), pos_h[:, :, 1::2].cos()), dim=3).flatten(2)

pos = torch.cat((pos_y, pos_x, pos_w, pos_h), dim=2)

return pos

@staticmethod

@compile_compatible_method_lru_cache(maxsize=4)

def build_sine_position_embedding(

shape: torch.Size,

device: torch.device | str,

dtype: torch.dtype,

num_position_features: int,

normalize: bool = False,

scale: float | None = None,

temperature: int = 10000,

mask: torch.Tensor | None = None,

) -> torch.Tensor:

batch_size, _, height, width = shape

if mask is None:

# Without a mask this is just a cumsum over ones, written out as arange

# instead: inductor's cumsum(ones) rewrite drops the requested dtype

# (https://github.com/pytorch/pytorch/issues/189518), which breaks

# float16/bfloat16 under torch.compile — don't revert to cumsum here

# until that fix is widely released.

y_embed = torch.arange(1, height + 1, dtype=dtype, device=device)[None, :, None].expand(

batch_size, height, width

)

x_embed = torch.arange(1, width + 1, dtype=dtype, device=device)[None, None, :].expand(

batch_size, height, width

)

else:

embed_mask = mask.to(dtype)

y_embed = embed_mask.cumsum(1)

x_embed = embed_mask.cumsum(2)

if normalize:

eps = 1e-6

y_embed = y_embed / (y_embed[:, -1:, :] + eps) * scale

x_embed = x_embed / (x_embed[:, :, -1:] + eps) * scale

dim_t = torch.arange(num_position_features, dtype=torch.int64, device=device).to(dtype)

dim_t = temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / num_position_features)

pos_x = x_embed[:, :, :, None] / dim_t

pos_y = y_embed[:, :, :, None] / dim_t

pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)

pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)

pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)

return pos

def forward(

self,

shape: torch.Size,

device: torch.device | str,

dtype: torch.dtype,

mask: torch.Tensor | None = None,

) -> torch.Tensor:

return self.build_sine_position_embedding(

shape, device, dtype, self.num_position_features, self.normalize, self.scale, self.temperature, mask

)

class Sam3FPNLayer(nn.Module):

def __init__(self, in_channels: int, fpn_dim: int, scale_factor: float):

super().__init__()

self.scale_factor = scale_factor

# Build the upsampling/downsampling layers based on scale factor

self.scale_layers = nn.ModuleList()

if scale_factor == 4.0:

self.scale_layers.append(nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2))

self.scale_layers.append(nn.GELU())

self.scale_layers.append(nn.ConvTranspose2d(in_channels // 2, in_channels // 4, kernel_size=2, stride=2))

intermediate_channels = in_channels // 4

elif scale_factor == 2.0:

self.scale_layers.append(nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2))

intermediate_channels = in_channels // 2