dgenerate-ultralytics-headless 8.3.214__py3-none-any.whl → 8.3.248__py3-none-any.whl

ultralytics/models/sam/modules/blocks.py

@@ -17,8 +17,7 @@ from .utils import add_decomposed_rel_pos, apply_rotary_enc, compute_axial_cis,
 
 
 class DropPath(nn.Module):
-    """
-    Implements stochastic depth regularization for neural networks during training.
+    """Implements stochastic depth regularization for neural networks during training.
 
     Attributes:
         drop_prob (float): Probability of dropping a path during training.
@@ -52,16 +51,14 @@ class DropPath(nn.Module):
 
 
 class MaskDownSampler(nn.Module):
-    """
-    A mask downsampling and embedding module for efficient processing of input masks.
+    """A mask downsampling and embedding module for efficient processing of input masks.
 
-    This class implements a mask downsampler that progressively reduces the spatial dimensions of input masks
-    while expanding their channel dimensions using convolutional layers, layer normalization, and activation
-    functions.
+    This class implements a mask downsampler that progressively reduces the spatial dimensions of input masks while
+    expanding their channel dimensions using convolutional layers, layer normalization, and activation functions.
 
     Attributes:
-        encoder (nn.Sequential): A sequential container of convolutional layers, layer normalization, and
-            activation functions for downsampling and embedding masks.
+        encoder (nn.Sequential): A sequential container of convolutional layers, layer normalization, and activation
+            functions for downsampling and embedding masks.
 
     Methods:
         forward: Downsamples and encodes input mask to embed_dim channels.
@@ -82,6 +79,7 @@ class MaskDownSampler(nn.Module):
         padding: int = 0,
         total_stride: int = 16,
         activation: type[nn.Module] = nn.GELU,
+        interpol_size: tuple[int, int] | None = None,
     ):
         """Initialize a mask downsampler module for progressive downsampling and channel expansion."""
         super().__init__()
@@ -105,18 +103,32 @@ class MaskDownSampler(nn.Module):
             mask_in_chans = mask_out_chans
 
         self.encoder.append(nn.Conv2d(mask_out_chans, embed_dim, kernel_size=1))
+        self.interpol_size = interpol_size
+        if self.interpol_size is not None:
+            assert isinstance(self.interpol_size, (list, tuple)), (
+                f"Unsupported type {type(self.interpol_size)}. Should be a list or tuple."
+            )
+            self.interpol_size = list(interpol_size)
+            assert len(self.interpol_size) == 2
 
     def forward(self, x: Tensor) -> Tensor:
         """Downsample and encode input mask to embed_dim channels using convolutional layers and LayerNorm2d."""
+        if self.interpol_size is not None and self.interpol_size != list(x.shape[-2:]):
+            x = F.interpolate(
+                x.float(),
+                size=self.interpol_size,
+                align_corners=False,
+                mode="bilinear",
+                antialias=True,
+            ).to(x.dtype)
         return self.encoder(x)
 
 
 class CXBlock(nn.Module):
-    """
-    ConvNeXt Block for efficient feature extraction in convolutional neural networks.
+    """ConvNeXt Block for efficient feature extraction in convolutional neural networks.
 
-    This block implements a modified version of the ConvNeXt architecture, offering improved performance and
-    flexibility in feature extraction.
+    This block implements a modified version of the ConvNeXt architecture, offering improved performance and flexibility
+    in feature extraction.
 
     Attributes:
         dwconv (nn.Conv2d): Depthwise or standard 2D convolution layer.
@@ -148,8 +160,7 @@ class CXBlock(nn.Module):
         layer_scale_init_value: float = 1e-6,
         use_dwconv: bool = True,
     ):
-        """
-        Initialize a ConvNeXt Block for efficient feature extraction in convolutional neural networks.
+        """Initialize a ConvNeXt Block for efficient feature extraction in convolutional neural networks.
 
         This block implements a modified version of the ConvNeXt architecture, offering improved performance and
         flexibility in feature extraction.
@@ -161,13 +172,6 @@ class CXBlock(nn.Module):
             drop_path (float): Stochastic depth rate.
             layer_scale_init_value (float): Initial value for Layer Scale.
             use_dwconv (bool): Whether to use depthwise convolution.
-
-        Examples:
-            >>> block = CXBlock(dim=64, kernel_size=7, padding=3)
-            >>> x = torch.randn(1, 64, 32, 32)
-            >>> output = block(x)
-            >>> print(output.shape)
-            torch.Size([1, 64, 32, 32])
         """
         super().__init__()
         self.dwconv = nn.Conv2d(
@@ -206,8 +210,7 @@ class CXBlock(nn.Module):
 
 
 class Fuser(nn.Module):
-    """
-    A module for fusing features through multiple layers of a neural network.
+    """A module for fusing features through multiple layers of a neural network.
 
     This class applies a series of identical layers to an input tensor, optionally projecting the input first.
 
@@ -228,8 +231,7 @@ class Fuser(nn.Module):
     """
 
     def __init__(self, layer: nn.Module, num_layers: int, dim: int | None = None, input_projection: bool = False):
-        """
-        Initialize the Fuser module for feature fusion through multiple layers.
+        """Initialize the Fuser module for feature fusion through multiple layers.
 
         This module creates a sequence of identical layers and optionally applies an input projection.
 
@@ -238,12 +240,6 @@ class Fuser(nn.Module):
             num_layers (int): The number of times to replicate the layer.
             dim (int | None): The dimension for input projection, if used.
             input_projection (bool): Whether to use input projection.
-
-        Examples:
-            >>> layer = nn.Linear(64, 64)
-            >>> fuser = Fuser(layer, num_layers=3, dim=64, input_projection=True)
-            >>> input_tensor = torch.randn(1, 64)
-            >>> output = fuser(input_tensor)
         """
         super().__init__()
         self.proj = nn.Identity()
@@ -262,12 +258,11 @@ class Fuser(nn.Module):
 
 
 class SAM2TwoWayAttentionBlock(TwoWayAttentionBlock):
-    """
-    A two-way attention block for performing self-attention and cross-attention in both directions.
+    """A two-way attention block for performing self-attention and cross-attention in both directions.
 
-    This block extends the TwoWayAttentionBlock and consists of four main components: self-attention on
-    sparse inputs, cross-attention from sparse to dense inputs, an MLP block on sparse inputs, and
-    cross-attention from dense to sparse inputs.
+    This block extends the TwoWayAttentionBlock and consists of four main components: self-attention on sparse inputs,
+    cross-attention from sparse to dense inputs, an MLP block on sparse inputs, and cross-attention from dense to sparse
+    inputs.
 
     Attributes:
         self_attn (Attention): Self-attention layer for queries.
@@ -299,12 +294,11 @@ class SAM2TwoWayAttentionBlock(TwoWayAttentionBlock):
         attention_downsample_rate: int = 2,
         skip_first_layer_pe: bool = False,
     ) -> None:
-        """
-        Initialize a SAM2TwoWayAttentionBlock for performing self-attention and cross-attention in two directions.
+        """Initialize a SAM2TwoWayAttentionBlock for performing self-attention and cross-attention in two directions.
 
         This block extends the TwoWayAttentionBlock and consists of four main components: self-attention on sparse
-        inputs, cross-attention from sparse to dense inputs, an MLP block on sparse inputs, and cross-attention
-        from dense to sparse inputs.
+        inputs, cross-attention from sparse to dense inputs, an MLP block on sparse inputs, and cross-attention from
+        dense to sparse inputs.
 
         Args:
             embedding_dim (int): The channel dimension of the embeddings.
@@ -313,24 +307,17 @@ class SAM2TwoWayAttentionBlock(TwoWayAttentionBlock):
             activation (Type[nn.Module]): The activation function of the MLP block.
             attention_downsample_rate (int): The downsample rate for attention computations.
             skip_first_layer_pe (bool): Whether to skip the positional encoding in the first layer.
-
-        Examples:
-            >>> block = SAM2TwoWayAttentionBlock(embedding_dim=256, num_heads=8, mlp_dim=2048)
-            >>> sparse_inputs = torch.randn(1, 100, 256)
-            >>> dense_inputs = torch.randn(1, 256, 32, 32)
-            >>> sparse_outputs, dense_outputs = block(sparse_inputs, dense_inputs)
         """
         super().__init__(embedding_dim, num_heads, mlp_dim, activation, attention_downsample_rate, skip_first_layer_pe)
         self.mlp = MLP(embedding_dim, mlp_dim, embedding_dim, num_layers=2, act=activation)
 
 
 class SAM2TwoWayTransformer(TwoWayTransformer):
-    """
-    A Two-Way Transformer module for simultaneous attention to image and query points.
+    """A Two-Way Transformer module for simultaneous attention to image and query points.
 
-    This class extends the TwoWayTransformer, implementing a specialized transformer decoder that attends to an
-    input image using queries with supplied positional embeddings. It is particularly useful for tasks like
-    object detection, image segmentation, and point cloud processing.
+    This class extends the TwoWayTransformer, implementing a specialized transformer decoder that attends to an input
+    image using queries with supplied positional embeddings. It is particularly useful for tasks like object detection,
+    image segmentation, and point cloud processing.
 
     Attributes:
         depth (int): Number of layers in the transformer.
@@ -362,11 +349,10 @@ class SAM2TwoWayTransformer(TwoWayTransformer):
         activation: type[nn.Module] = nn.ReLU,
         attention_downsample_rate: int = 2,
     ) -> None:
-        """
-        Initialize a SAM2TwoWayTransformer instance.
+        """Initialize a SAM2TwoWayTransformer instance.
 
-        This transformer decoder attends to an input image using queries with supplied positional embeddings.
-        It is designed for tasks like object detection, image segmentation, and point cloud processing.
+        This transformer decoder attends to an input image using queries with supplied positional embeddings. It is
+        designed for tasks like object detection, image segmentation, and point cloud processing.
 
         Args:
             depth (int): Number of layers in the transformer.
@@ -375,17 +361,6 @@ class SAM2TwoWayTransformer(TwoWayTransformer):
             mlp_dim (int): Channel dimension internal to the MLP block.
             activation (Type[nn.Module]): Activation function to use in the MLP block.
             attention_downsample_rate (int): Downsampling rate for attention computations.
-
-        Examples:
-            >>> transformer = SAM2TwoWayTransformer(depth=5, embedding_dim=256, num_heads=8, mlp_dim=2048)
-            >>> transformer
-            SAM2TwoWayTransformer(
-              (layers): ModuleList(
-                (0-4): 5 x SAM2TwoWayAttentionBlock(...)
-              )
-              (final_attn_token_to_image): Attention(...)
-              (norm_final_attn): LayerNorm(...)
-            )
         """
         super().__init__(depth, embedding_dim, num_heads, mlp_dim, activation, attention_downsample_rate)
         self.layers = nn.ModuleList()
@@ -403,11 +378,10 @@ class SAM2TwoWayTransformer(TwoWayTransformer):
 
 
 class RoPEAttention(Attention):
-    """
-    Implements rotary position encoding for attention mechanisms in transformer architectures.
+    """Implements rotary position encoding for attention mechanisms in transformer architectures.
 
-    This class extends the base Attention class by incorporating Rotary Position Encoding (RoPE) to enhance
-    the positional awareness of the attention mechanism.
+    This class extends the base Attention class by incorporating Rotary Position Encoding (RoPE) to enhance the
+    positional awareness of the attention mechanism.
 
     Attributes:
         compute_cis (Callable): Function to compute axial complex numbers for rotary encoding.
@@ -471,13 +445,7 @@ class RoPEAttention(Attention):
         )
 
         # Attention
-        _, _, _, c_per_head = q.shape
-        attn = q @ k.permute(0, 1, 3, 2)  # B x N_heads x N_tokens x N_tokens
-        attn = attn / math.sqrt(c_per_head)
-        attn = torch.softmax(attn, dim=-1)
-
-        # Get output
-        out = attn @ v
+        out = F.scaled_dot_product_attention(q, k, v)
 
         out = self._recombine_heads(out)
         out = self.out_proj(out)
@@ -501,12 +469,11 @@ def do_pool(x: torch.Tensor, pool: nn.Module, norm: nn.Module = None) -> torch.T
 
 
 class MultiScaleAttention(nn.Module):
-    """
-    Implements multiscale self-attention with optional query pooling for efficient feature extraction.
+    """Implements multiscale self-attention with optional query pooling for efficient feature extraction.
 
-    This class provides a flexible implementation of multiscale attention, allowing for optional
-    downsampling of query features through pooling. It's designed to enhance the model's ability to
-    capture multiscale information in visual tasks.
+    This class provides a flexible implementation of multiscale attention, allowing for optional downsampling of query
+    features through pooling. It's designed to enhance the model's ability to capture multiscale information in visual
+    tasks.
 
     Attributes:
         dim (int): Input dimension of the feature map.
@@ -581,11 +548,10 @@ class MultiScaleAttention(nn.Module):
 
 
 class MultiScaleBlock(nn.Module):
-    """
-    A multiscale attention block with window partitioning and query pooling for efficient vision transformers.
+    """A multiscale attention block with window partitioning and query pooling for efficient vision transformers.
 
-    This class implements a multiscale attention mechanism with optional window partitioning and downsampling,
-    designed for use in vision transformer architectures.
+    This class implements a multiscale attention mechanism with optional window partitioning and downsampling, designed
+    for use in vision transformer architectures.
 
     Attributes:
         dim (int): Input dimension of the block.
@@ -619,7 +585,7 @@ class MultiScaleBlock(nn.Module):
         mlp_ratio: float = 4.0,
         drop_path: float = 0.0,
         norm_layer: nn.Module | str = "LayerNorm",
-        q_stride: tuple[int, int] = None,
+        q_stride: tuple[int, int] | None = None,
         act_layer: type[nn.Module] = nn.GELU,
         window_size: int = 0,
     ):
@@ -696,11 +662,10 @@ class MultiScaleBlock(nn.Module):
 
 
 class PositionEmbeddingSine(nn.Module):
-    """
-    A module for generating sinusoidal positional embeddings for 2D inputs like images.
+    """A module for generating sinusoidal positional embeddings for 2D inputs like images.
 
-    This class implements sinusoidal position encoding for 2D spatial positions, which can be used in
-    transformer-based models for computer vision tasks.
+    This class implements sinusoidal position encoding for 2D spatial positions, which can be used in transformer-based
+    models for computer vision tasks.
 
     Attributes:
         num_pos_feats (int): Number of positional features (half of the embedding dimension).
@@ -811,8 +776,7 @@ class PositionEmbeddingSine(nn.Module):
 
 
 class PositionEmbeddingRandom(nn.Module):
-    """
-    Positional encoding using random spatial frequencies.
+    """Positional encoding using random spatial frequencies.
 
     This class generates positional embeddings for input coordinates using random spatial frequencies. It is
     particularly useful for transformer-based models that require position information.
@@ -878,12 +842,11 @@ class PositionEmbeddingRandom(nn.Module):
 
 
 class Block(nn.Module):
-    """
-    Transformer block with support for window attention and residual propagation.
+    """Transformer block with support for window attention and residual propagation.
 
-    This class implements a transformer block that can use either global or windowed self-attention,
-    followed by a feed-forward network. It supports relative positional embeddings and is designed
-    for use in vision transformer architectures.
+    This class implements a transformer block that can use either global or windowed self-attention, followed by a
+    feed-forward network. It supports relative positional embeddings and is designed for use in vision transformer
+    architectures.
 
     Attributes:
         norm1 (nn.Module): First normalization layer.
@@ -917,12 +880,11 @@ class Block(nn.Module):
         window_size: int = 0,
         input_size: tuple[int, int] | None = None,
     ) -> None:
-        """
-        Initialize a transformer block with optional window attention and relative positional embeddings.
+        """Initialize a transformer block with optional window attention and relative positional embeddings.
 
-        This constructor sets up a transformer block that can use either global or windowed self-attention,
-        followed by a feed-forward network. It supports relative positional embeddings and is designed
-        for use in vision transformer architectures.
+        This constructor sets up a transformer block that can use either global or windowed self-attention, followed by
+        a feed-forward network. It supports relative positional embeddings and is designed for use in vision transformer
+        architectures.
 
         Args:
             dim (int): Number of input channels.
@@ -935,13 +897,6 @@ class Block(nn.Module):
             rel_pos_zero_init (bool): If True, initializes relative positional parameters to zero.
             window_size (int): Size of attention window. If 0, uses global attention.
             input_size (tuple[int, int] | None): Input resolution for calculating relative positional parameter size.
-
-        Examples:
-            >>> block = Block(dim=256, num_heads=8, window_size=7)
-            >>> x = torch.randn(1, 56, 56, 256)
-            >>> output = block(x)
-            >>> print(output.shape)
-            torch.Size([1, 56, 56, 256])
         """
         super().__init__()
         self.norm1 = norm_layer(dim)
@@ -978,12 +933,11 @@ class Block(nn.Module):
 
 
 class REAttention(nn.Module):
-    """
-    Relative Position Attention module for efficient self-attention in transformer architectures.
+    """Relative Position Attention module for efficient self-attention in transformer architectures.
 
-    This class implements a multi-head attention mechanism with relative positional embeddings, designed
-    for use in vision transformer models. It supports optional query pooling and window partitioning
-    for efficient processing of large inputs.
+    This class implements a multi-head attention mechanism with relative positional embeddings, designed for use in
+    vision transformer models. It supports optional query pooling and window partitioning for efficient processing of
+    large inputs.
 
     Attributes:
         num_heads (int): Number of attention heads.
@@ -1014,11 +968,10 @@ class REAttention(nn.Module):
         rel_pos_zero_init: bool = True,
         input_size: tuple[int, int] | None = None,
     ) -> None:
-        """
-        Initialize a Relative Position Attention module for transformer-based architectures.
+        """Initialize a Relative Position Attention module for transformer-based architectures.
 
-        This module implements multi-head attention with optional relative positional encodings, designed
-        specifically for vision tasks in transformer models.
+        This module implements multi-head attention with optional relative positional encodings, designed specifically
+        for vision tasks in transformer models.
 
         Args:
             dim (int): Number of input channels.
@@ -1028,13 +981,6 @@ class REAttention(nn.Module):
             rel_pos_zero_init (bool): If True, initializes relative positional parameters to zero.
             input_size (tuple[int, int] | None): Input resolution for calculating relative positional parameter size.
                 Required if use_rel_pos is True.
-
-        Examples:
-            >>> attention = REAttention(dim=256, num_heads=8, input_size=(32, 32))
-            >>> x = torch.randn(1, 32, 32, 256)
-            >>> output = attention(x)
-            >>> print(output.shape)
-            torch.Size([1, 32, 32, 256])
         """
         super().__init__()
         self.num_heads = num_heads
@@ -1070,12 +1016,11 @@ class REAttention(nn.Module):
 
 
 class PatchEmbed(nn.Module):
-    """
-    Image to Patch Embedding module for vision transformer architectures.
+    """Image to Patch Embedding module for vision transformer architectures.
 
-    This module converts an input image into a sequence of patch embeddings using a convolutional layer.
-    It is commonly used as the first layer in vision transformer architectures to transform image data
-    into a suitable format for subsequent transformer blocks.
+    This module converts an input image into a sequence of patch embeddings using a convolutional layer. It is commonly
+    used as the first layer in vision transformer architectures to transform image data into a suitable format for
+    subsequent transformer blocks.
 
     Attributes:
         proj (nn.Conv2d): Convolutional layer for projecting image patches to embeddings.
@@ -1098,12 +1043,12 @@ class PatchEmbed(nn.Module):
         padding: tuple[int, int] = (0, 0),
         in_chans: int = 3,
         embed_dim: int = 768,
+        bias: bool = True,
     ) -> None:
-        """
-        Initialize the PatchEmbed module for converting image patches to embeddings.
+        """Initialize the PatchEmbed module for converting image patches to embeddings.
 
-        This module is typically used as the first layer in vision transformer architectures to transform
-        image data into a suitable format for subsequent transformer blocks.
+        This module is typically used as the first layer in vision transformer architectures to transform image data
+        into a suitable format for subsequent transformer blocks.
 
         Args:
             kernel_size (tuple[int, int]): Size of the convolutional kernel for patch extraction.
@@ -1111,17 +1056,11 @@ class PatchEmbed(nn.Module):
             padding (tuple[int, int]): Padding applied to the input before convolution.
             in_chans (int): Number of input image channels.
             embed_dim (int): Dimensionality of the output patch embeddings.
-
-        Examples:
-            >>> patch_embed = PatchEmbed(kernel_size=(16, 16), stride=(16, 16), in_chans=3, embed_dim=768)
-            >>> x = torch.randn(1, 3, 224, 224)
-            >>> output = patch_embed(x)
-            >>> print(output.shape)
-            torch.Size([1, 768, 14, 14])
+            bias (bool): Whether to include a bias term in the convolutional layer.
         """
         super().__init__()
 
-        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding)
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Compute patch embedding by applying convolution and transposing resulting tensor."""

ultralytics/models/sam/modules/decoders.py

@@ -9,8 +9,7 @@ from ultralytics.nn.modules import MLP, LayerNorm2d
 
 
 class MaskDecoder(nn.Module):
-    """
-    Decoder module for generating masks and their associated quality scores using a transformer architecture.
+    """Decoder module for generating masks and their associated quality scores using a transformer architecture.
 
     This class predicts masks given image and prompt embeddings, utilizing a transformer to process the inputs and
     generate mask predictions along with their quality scores.
@@ -47,8 +46,7 @@ class MaskDecoder(nn.Module):
         iou_head_depth: int = 3,
         iou_head_hidden_dim: int = 256,
     ) -> None:
-        """
-        Initialize the MaskDecoder module for generating masks and their associated quality scores.
+        """Initialize the MaskDecoder module for generating masks and their associated quality scores.
 
         Args:
             transformer_dim (int): Channel dimension for the transformer module.
@@ -57,11 +55,6 @@ class MaskDecoder(nn.Module):
             activation (Type[nn.Module]): Type of activation to use when upscaling masks.
             iou_head_depth (int): Depth of the MLP used to predict mask quality.
             iou_head_hidden_dim (int): Hidden dimension of the MLP used to predict mask quality.
-
-        Examples:
-            >>> transformer = nn.TransformerEncoder(nn.TransformerEncoderLayer(d_model=256, nhead=8), num_layers=6)
-            >>> decoder = MaskDecoder(transformer_dim=256, transformer=transformer)
-            >>> print(decoder)
         """
         super().__init__()
         self.transformer_dim = transformer_dim
@@ -94,8 +87,7 @@ class MaskDecoder(nn.Module):
         dense_prompt_embeddings: torch.Tensor,
         multimask_output: bool,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Predict masks given image and prompt embeddings.
+        """Predict masks given image and prompt embeddings.
 
         Args:
             image_embeddings (torch.Tensor): Embeddings from the image encoder.
@@ -172,11 +164,10 @@ class MaskDecoder(nn.Module):
 
 
 class SAM2MaskDecoder(nn.Module):
-    """
-    Transformer-based decoder for predicting instance segmentation masks from image and prompt embeddings.
+    """Transformer-based decoder for predicting instance segmentation masks from image and prompt embeddings.
 
-    This class extends the functionality of the MaskDecoder, incorporating additional features such as
-    high-resolution feature processing, dynamic multimask output, and object score prediction.
+    This class extends the functionality of the MaskDecoder, incorporating additional features such as high-resolution
+    feature processing, dynamic multimask output, and object score prediction.
 
     Attributes:
         transformer_dim (int): Channel dimension of the transformer.
@@ -233,11 +224,10 @@ class SAM2MaskDecoder(nn.Module):
         pred_obj_scores_mlp: bool = False,
         use_multimask_token_for_obj_ptr: bool = False,
     ) -> None:
-        """
-        Initialize the SAM2MaskDecoder module for predicting instance segmentation masks.
+        """Initialize the SAM2MaskDecoder module for predicting instance segmentation masks.
 
-        This decoder extends the functionality of MaskDecoder, incorporating additional features such as
-        high-resolution feature processing, dynamic multimask output, and object score prediction.
+        This decoder extends the functionality of MaskDecoder, incorporating additional features such as high-resolution
+        feature processing, dynamic multimask output, and object score prediction.
 
         Args:
             transformer_dim (int): Channel dimension of the transformer.
@@ -254,11 +244,6 @@ class SAM2MaskDecoder(nn.Module):
             pred_obj_scores (bool): Whether to predict object scores.
             pred_obj_scores_mlp (bool): Whether to use MLP for object score prediction.
             use_multimask_token_for_obj_ptr (bool): Whether to use multimask token for object pointer.
-
-        Examples:
-            >>> transformer = nn.TransformerEncoder(nn.TransformerEncoderLayer(d_model=256, nhead=8), num_layers=6)
-            >>> decoder = SAM2MaskDecoder(transformer_dim=256, transformer=transformer)
-            >>> print(decoder)
         """
         super().__init__()
         self.transformer_dim = transformer_dim
@@ -319,8 +304,7 @@ class SAM2MaskDecoder(nn.Module):
         repeat_image: bool,
         high_res_features: list[torch.Tensor] | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """
-        Predict masks given image and prompt embeddings.
+        """Predict masks given image and prompt embeddings.
 
         Args:
             image_embeddings (torch.Tensor): Embeddings from the image encoder with shape (B, C, H, W).
@@ -452,23 +436,21 @@ class SAM2MaskDecoder(nn.Module):
     def _get_stability_scores(self, mask_logits):
         """Compute mask stability scores based on IoU between upper and lower thresholds."""
         mask_logits = mask_logits.flatten(-2)
-        stability_delta = self.dynamic_multimask_stability_delta
-        area_i = torch.sum(mask_logits > stability_delta, dim=-1).float()
-        area_u = torch.sum(mask_logits > -stability_delta, dim=-1).float()
+        area_i = torch.sum(mask_logits > self.dynamic_multimask_stability_delta, dim=-1).float()
+        area_u = torch.sum(mask_logits > -self.dynamic_multimask_stability_delta, dim=-1).float()
         return torch.where(area_u > 0, area_i / area_u, 1.0)
 
     def _dynamic_multimask_via_stability(self, all_mask_logits, all_iou_scores):
-        """
-        Dynamically select the most stable mask output based on stability scores and IoU predictions.
+        """Dynamically select the most stable mask output based on stability scores and IoU predictions.
 
-        This method is used when outputting a single mask. If the stability score from the current single-mask
-        output (based on output token 0) falls below a threshold, it instead selects from multi-mask outputs
-        (based on output tokens 1-3) the mask with the highest predicted IoU score. This ensures a valid mask
-        for both clicking and tracking scenarios.
+        This method is used when outputting a single mask. If the stability score from the current single-mask output
+        (based on output token 0) falls below a threshold, it instead selects from multi-mask outputs (based on output
+        tokens 1-3) the mask with the highest predicted IoU score. This ensures a valid mask for both clicking and
+        tracking scenarios.
 
         Args:
-            all_mask_logits (torch.Tensor): Logits for all predicted masks, shape (B, N, H, W) where B is
-                batch size, N is number of masks (typically 4), and H, W are mask dimensions.
+            all_mask_logits (torch.Tensor): Logits for all predicted masks, shape (B, N, H, W) where B is batch size, N
+                is number of masks (typically 4), and H, W are mask dimensions.
             all_iou_scores (torch.Tensor): Predicted IoU scores for all masks, shape (B, N).
 
         Returns: