ultralytics 8.3.101__py3-none-any.whl → 8.3.102__py3-none-any.whl

ultralytics/models/yolo/yoloe/train.py

@@ -38,7 +38,23 @@ class YOLOETrainer(DetectionTrainer):
         super().__init__(cfg, overrides, _callbacks)
 
     def get_model(self, cfg=None, weights=None, verbose=True):
-        """Return YOLOEModel initialized with specified config and weights."""
+        """
+        Return a YOLOEModel initialized with the specified configuration and weights.
+
+        Args:
+            cfg (dict | str | None): Model configuration. Can be a dictionary containing a 'yaml_file' key,
+                a direct path to a YAML file, or None to use default configuration.
+            weights (str | Path | None): Path to pretrained weights file to load into the model.
+            verbose (bool): Whether to display model information during initialization.
+
+        Returns:
+            (YOLOEModel): The initialized YOLOE model.
+
+        Notes:
+            - The number of classes (nc) is hard-coded to a maximum of 80 following the official configuration.
+            - The nc parameter here represents the maximum number of different text samples in one image,
+              rather than the actual number of classes.
+        """
         # NOTE: This `nc` here is the max number of different text samples in one image, rather than the actual `nc`.
         # NOTE: Following the official config, nc hard-coded to 80 for now.
         model = YOLOEModel(
@@ -180,7 +196,20 @@ class YOLOETrainerFromScratch(YOLOETrainer):
         return YOLOConcatDataset(datasets) if len(datasets) > 1 else datasets[0]
 
     def set_text_embeddings(self, datasets, batch):
-        """Set text embeddings for datasets to accelerate training by caching category names."""
+        """
+        Set text embeddings for datasets to accelerate training by caching category names.
+
+        This method collects unique category names from all datasets, then generates and caches text embeddings
+        for these categories to improve training efficiency.
+
+        Args:
+            datasets (List[Dataset]): List of datasets from which to extract category names.
+            batch (int | None): Batch size used for processing.
+
+        Notes:
+            This method collects category names from datasets that have the 'category_names' attribute,
+            then uses the first dataset's image path to determine where to cache the generated text embeddings.
+        """
         # TODO: open up an interface to determine whether to do cache
         category_names = set()
         for dataset in datasets:
@@ -312,7 +341,22 @@ class YOLOEPEFreeTrainer(YOLOEPETrainer, YOLOETrainerFromScratch):
         return batch
 
     def set_text_embeddings(self, datasets, batch):
-        """No need to set text embeddings for prompt-free fine-tuning."""
+        """
+        Set text embeddings for datasets to accelerate training by caching category names.
+
+        This method collects unique category names from all datasets, generates text embeddings for them,
+        and caches these embeddings to improve training efficiency. The embeddings are stored in a file
+        in the parent directory of the first dataset's image path.
+
+        Args:
+            datasets (List[Dataset]): List of datasets containing category names to process.
+            batch (int): Batch size for processing text embeddings.
+
+        Notes:
+            The method creates a dictionary mapping text samples to their embeddings and stores it
+            at the path specified by 'cache_path'. If the cache file already exists, it will be loaded
+            instead of regenerating the embeddings.
+        """
         pass
 
 

ultralytics/nn/modules/activation.py

@@ -7,14 +7,26 @@ import torch.nn as nn
 
 class AGLU(nn.Module):
     """
-    Unified activation function module from https://github.com/kostas1515/AGLU.
+    Unified activation function module from AGLU.
 
-    This class implements a parameterized activation function with learnable parameters lambda and kappa.
+    This class implements a parameterized activation function with learnable parameters lambda and kappa, based on the
+    AGLU (Adaptive Gated Linear Unit) approach (https://github.com/kostas1515/AGLU).
 
     Attributes:
         act (nn.Softplus): Softplus activation function with negative beta.
         lambd (nn.Parameter): Learnable lambda parameter initialized with uniform distribution.
         kappa (nn.Parameter): Learnable kappa parameter initialized with uniform distribution.
+
+    Methods:
+        forward: Compute the forward pass of the Unified activation function.
+
+    Examples:
+        >>> import torch
+        >>> m = AGLU()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+        >>> print(output.shape)
+        torch.Size([2])
     """
 
     def __init__(self, device=None, dtype=None) -> None:
@@ -25,6 +37,17 @@ class AGLU(nn.Module):
         self.kappa = nn.Parameter(nn.init.uniform_(torch.empty(1, device=device, dtype=dtype)))  # kappa parameter
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Compute the forward pass of the Unified activation function."""
+        """
+        Apply the Adaptive Gated Linear Unit (AGLU) activation function.
+
+        This forward method implements the AGLU activation function with learnable parameters lambda and kappa.
+        The function applies a transformation that adaptively combines linear and non-linear components.
+
+        Args:
+            x (torch.Tensor): Input tensor to apply the activation function to.
+
+        Returns:
+            (torch.Tensor): Output tensor after applying the AGLU activation function, with the same shape as the input.
+        """
         lam = torch.clamp(self.lambd, min=0.0001)  # Clamp lambda to avoid division by zero
         return torch.exp((1 / lam) * self.act((self.kappa * x) - torch.log(lam)))

ultralytics/nn/modules/block.py

@@ -1875,3 +1875,92 @@ class A2C2f(nn.Module):
         if self.gamma is not None:
             return x + self.gamma.view(-1, len(self.gamma), 1, 1) * y
         return y
+
+
+class SwiGLUFFN(nn.Module):
+    """SwiGLU Feed-Forward Network for transformer-based architectures."""
+
+    def __init__(self, gc, ec, e=4) -> None:
+        """Initialize SwiGLU FFN with input dimension, output dimension, and expansion factor."""
+        super().__init__()
+        self.w12 = nn.Linear(gc, e * ec)
+        self.w3 = nn.Linear(e * ec // 2, ec)
+
+    def forward(self, x):
+        """Apply SwiGLU transformation to input features."""
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return self.w3(hidden)
+
+
+class Residual(nn.Module):
+    """Residual connection wrapper for neural network modules."""
+
+    def __init__(self, m) -> None:
+        """Initialize residual module with the wrapped module."""
+        super().__init__()
+        self.m = m
+        nn.init.zeros_(self.m.w3.bias)
+        # For models with l scale, please change the initialization to
+        # nn.init.constant_(self.m.w3.weight, 1e-6)
+        nn.init.zeros_(self.m.w3.weight)
+
+    def forward(self, x):
+        """Apply residual connection to input features."""
+        return x + self.m(x)
+
+
+class SAVPE(nn.Module):
+    """Spatial-Aware Visual Prompt Embedding module for feature enhancement."""
+
+    def __init__(self, ch, c3, embed):
+        """Initialize SAVPE module with channels, intermediate channels, and embedding dimension."""
+        super().__init__()
+        self.cv1 = nn.ModuleList(
+            nn.Sequential(
+                Conv(x, c3, 3), Conv(c3, c3, 3), nn.Upsample(scale_factor=i * 2) if i in {1, 2} else nn.Identity()
+            )
+            for i, x in enumerate(ch)
+        )
+
+        self.cv2 = nn.ModuleList(
+            nn.Sequential(Conv(x, c3, 1), nn.Upsample(scale_factor=i * 2) if i in {1, 2} else nn.Identity())
+            for i, x in enumerate(ch)
+        )
+
+        self.c = 16
+        self.cv3 = nn.Conv2d(3 * c3, embed, 1)
+        self.cv4 = nn.Conv2d(3 * c3, self.c, 3, padding=1)
+        self.cv5 = nn.Conv2d(1, self.c, 3, padding=1)
+        self.cv6 = nn.Sequential(Conv(2 * self.c, self.c, 3), nn.Conv2d(self.c, self.c, 3, padding=1))
+
+    def forward(self, x, vp):
+        """Process input features and visual prompts to generate enhanced embeddings."""
+        y = [self.cv2[i](xi) for i, xi in enumerate(x)]
+        y = self.cv4(torch.cat(y, dim=1))
+
+        x = [self.cv1[i](xi) for i, xi in enumerate(x)]
+        x = self.cv3(torch.cat(x, dim=1))
+
+        B, C, H, W = x.shape
+
+        Q = vp.shape[1]
+
+        x = x.view(B, C, -1)
+
+        y = y.reshape(B, 1, self.c, H, W).expand(-1, Q, -1, -1, -1).reshape(B * Q, self.c, H, W)
+        vp = vp.reshape(B, Q, 1, H, W).reshape(B * Q, 1, H, W)
+
+        y = self.cv6(torch.cat((y, self.cv5(vp)), dim=1))
+
+        y = y.reshape(B, Q, self.c, -1)
+        vp = vp.reshape(B, Q, 1, -1)
+
+        score = y * vp + torch.logical_not(vp) * torch.finfo(y.dtype).min
+
+        score = F.softmax(score, dim=-1, dtype=torch.float).to(score.dtype)
+
+        aggregated = score.transpose(-2, -3) @ x.reshape(B, self.c, C // self.c, -1).transpose(-1, -2)
+
+        return F.normalize(aggregated.transpose(-2, -3).reshape(B, Q, -1), dim=-1, p=2)

ultralytics/nn/modules/head.py

@@ -12,7 +12,7 @@ from torch.nn.init import constant_, xavier_uniform_
 from ultralytics.utils.tal import TORCH_1_10, dist2bbox, dist2rbox, make_anchors
 from ultralytics.utils.torch_utils import fuse_conv_and_bn, smart_inference_mode
 
-from .block import DFL, BNContrastiveHead, ContrastiveHead, Proto
+from .block import DFL, SAVPE, BNContrastiveHead, ContrastiveHead, Proto, Residual, SwiGLUFFN
 from .conv import Conv, DWConv
 from .transformer import MLP, DeformableTransformerDecoder, DeformableTransformerDecoderLayer
 from .utils import bias_init_with_prob, linear_init
@@ -345,61 +345,6 @@ class WorldDetect(Detect):
             # b[-1].bias.data[:] = math.log(5 / m.nc / (640 / s) ** 2)  # cls (.01 objects, 80 classes, 640 img)
 
 
-class SAVPE(nn.Module):
-    """Spatial-Aware Visual Prompt Embedding module for feature enhancement."""
-
-    def __init__(self, ch, c3, embed):
-        """Initialize SAVPE module with channels, intermediate channels, and embedding dimension."""
-        super().__init__()
-        self.cv1 = nn.ModuleList(
-            nn.Sequential(
-                Conv(x, c3, 3), Conv(c3, c3, 3), nn.Upsample(scale_factor=i * 2) if i in {1, 2} else nn.Identity()
-            )
-            for i, x in enumerate(ch)
-        )
-
-        self.cv2 = nn.ModuleList(
-            nn.Sequential(Conv(x, c3, 1), nn.Upsample(scale_factor=i * 2) if i in {1, 2} else nn.Identity())
-            for i, x in enumerate(ch)
-        )
-
-        self.c = 16
-        self.cv3 = nn.Conv2d(3 * c3, embed, 1)
-        self.cv4 = nn.Conv2d(3 * c3, self.c, 3, padding=1)
-        self.cv5 = nn.Conv2d(1, self.c, 3, padding=1)
-        self.cv6 = nn.Sequential(Conv(2 * self.c, self.c, 3), nn.Conv2d(self.c, self.c, 3, padding=1))
-
-    def forward(self, x, vp):
-        """Process input features and visual prompts to generate enhanced embeddings."""
-        y = [self.cv2[i](xi) for i, xi in enumerate(x)]
-        y = self.cv4(torch.cat(y, dim=1))
-
-        x = [self.cv1[i](xi) for i, xi in enumerate(x)]
-        x = self.cv3(torch.cat(x, dim=1))
-
-        B, C, H, W = x.shape
-
-        Q = vp.shape[1]
-
-        x = x.view(B, C, -1)
-
-        y = y.reshape(B, 1, self.c, H, W).expand(-1, Q, -1, -1, -1).reshape(B * Q, self.c, H, W)
-        vp = vp.reshape(B, Q, 1, H, W).reshape(B * Q, 1, H, W)
-
-        y = self.cv6(torch.cat((y, self.cv5(vp)), dim=1))
-
-        y = y.reshape(B, Q, self.c, -1)
-        vp = vp.reshape(B, Q, 1, -1)
-
-        score = y * vp + torch.logical_not(vp) * torch.finfo(y.dtype).min
-
-        score = F.softmax(score, dim=-1, dtype=torch.float).to(score.dtype)
-
-        aggregated = score.transpose(-2, -3) @ x.reshape(B, self.c, C // self.c, -1).transpose(-1, -2)
-
-        return F.normalize(aggregated.transpose(-2, -3).reshape(B, Q, -1), dim=-1, p=2)
-
-
 class LRPCHead(nn.Module):
     """Lightweight Region Proposal and Classification Head for efficient object detection."""
 
@@ -419,7 +364,7 @@ class LRPCHead(nn.Module):
         linear.bias.data = conv.bias.data
         return linear
 
-    def forward(self, cls_feat, loc_feat, conf, max_det):
+    def forward(self, cls_feat, loc_feat, conf):
         """Process classification and localization features to generate detection proposals."""
         if self.enabled:
             pf_score = self.pf(cls_feat)[0, 0].flatten(0)
@@ -533,7 +478,7 @@ class YOLOEDetect(Detect):
             cls_feat = self.cv3[i](x[i])
             loc_feat = self.cv2[i](x[i])
             assert isinstance(self.lrpc[i], LRPCHead)
-            x[i], mask = self.lrpc[i](cls_feat, loc_feat, self.conf, self.max_det)
+            x[i], mask = self.lrpc[i](cls_feat, loc_feat, getattr(self, "conf", 0.001))
             masks.append(mask)
         shape = x[0][0].shape
         if self.dynamic or self.shape != shape:
@@ -585,40 +530,6 @@ class YOLOEDetect(Detect):
             c.bias.data[:] = math.log(5 / m.nc / (640 / s) ** 2)
 
 
-class SwiGLUFFN(nn.Module):
-    """SwiGLU Feed-Forward Network for transformer-based architectures."""
-
-    def __init__(self, gc, ec, e=4) -> None:
-        """Initialize SwiGLU FFN with input dimension, output dimension, and expansion factor."""
-        super().__init__()
-        self.w12 = nn.Linear(gc, e * ec)
-        self.w3 = nn.Linear(e * ec // 2, ec)
-
-    def forward(self, x):
-        """Apply SwiGLU transformation to input features."""
-        x12 = self.w12(x)
-        x1, x2 = x12.chunk(2, dim=-1)
-        hidden = F.silu(x1) * x2
-        return self.w3(hidden)
-
-
-class Residual(nn.Module):
-    """Residual connection wrapper for neural network modules."""
-
-    def __init__(self, m) -> None:
-        """Initialize residual module with the wrapped module."""
-        super().__init__()
-        self.m = m
-        nn.init.zeros_(self.m.w3.bias)
-        # For models with l scale, please change the initialization to
-        # nn.init.constant_(self.m.w3.weight, 1e-6)
-        nn.init.zeros_(self.m.w3.weight)
-
-    def forward(self, x):
-        """Apply residual connection to input features."""
-        return x + self.m(x)
-
-
 class YOLOESegment(YOLOEDetect):
     """YOLO segmentation head with text embedding capabilities."""
 

ultralytics/nn/modules/utils.py

@@ -13,17 +13,66 @@ __all__ = "multi_scale_deformable_attn_pytorch", "inverse_sigmoid"
 
 
 def _get_clones(module, n):
-    """Create a list of cloned modules from the given module."""
+    """
+    Create a list of cloned modules from the given module.
+
+    Args:
+        module (nn.Module): The module to be cloned.
+        n (int): Number of clones to create.
+
+    Returns:
+        (nn.ModuleList): A ModuleList containing n clones of the input module.
+
+    Examples:
+        >>> import torch.nn as nn
+        >>> layer = nn.Linear(10, 10)
+        >>> clones = _get_clones(layer, 3)
+        >>> len(clones)
+        3
+    """
     return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])
 
 
 def bias_init_with_prob(prior_prob=0.01):
-    """Initialize conv/fc bias value according to a given probability value."""
+    """
+    Initialize conv/fc bias value according to a given probability value.
+
+    This function calculates the bias initialization value based on a prior probability using the inverse error function.
+    It's commonly used in object detection models to initialize classification layers with a specific positive prediction
+    probability.
+
+    Args:
+        prior_prob (float, optional): Prior probability for bias initialization.
+
+    Returns:
+        (float): Bias initialization value calculated from the prior probability.
+
+    Examples:
+        >>> bias = bias_init_with_prob(0.01)
+        >>> print(f"Bias initialization value: {bias:.4f}")
+        Bias initialization value: -4.5951
+    """
     return float(-np.log((1 - prior_prob) / prior_prob))  # return bias_init
 
 
 def linear_init(module):
-    """Initialize the weights and biases of a linear module."""
+    """
+    Initialize the weights and biases of a linear module.
+
+    This function initializes the weights of a linear module using a uniform distribution within bounds calculated
+    from the input dimension. If the module has a bias, it is also initialized.
+
+    Args:
+        module (nn.Module): Linear module to initialize.
+
+    Returns:
+        (nn.Module): The initialized module.
+
+    Examples:
+        >>> import torch.nn as nn
+        >>> linear = nn.Linear(10, 5)
+        >>> initialized_linear = linear_init(linear)
+    """
     bound = 1 / math.sqrt(module.weight.shape[0])
     uniform_(module.weight, -bound, bound)
     if hasattr(module, "bias") and module.bias is not None:
@@ -31,7 +80,24 @@ def linear_init(module):
 
 
 def inverse_sigmoid(x, eps=1e-5):
-    """Calculate the inverse sigmoid function for a tensor."""
+    """
+    Calculate the inverse sigmoid function for a tensor.
+
+    This function applies the inverse of the sigmoid function to a tensor, which is useful in various neural network
+    operations, particularly in attention mechanisms and coordinate transformations.
+
+    Args:
+        x (torch.Tensor): Input tensor with values in range [0, 1].
+        eps (float, optional): Small epsilon value to prevent numerical instability.
+
+    Returns:
+        (torch.Tensor): Tensor after applying the inverse sigmoid function.
+
+    Examples:
+        >>> x = torch.tensor([0.2, 0.5, 0.8])
+        >>> inverse_sigmoid(x)
+        tensor([-1.3863,  0.0000,  1.3863])
+    """
     x = x.clamp(min=0, max=1)
     x1 = x.clamp(min=eps)
     x2 = (1 - x).clamp(min=eps)

ultralytics/nn/text_model.py

@@ -6,7 +6,7 @@ from pathlib import Path
 import torch
 import torch.nn as nn
 
-from ultralytics.utils import LOGGER, checks
+from ultralytics.utils import checks
 from ultralytics.utils.torch_utils import smart_inference_mode
 
 try:
@@ -59,9 +59,10 @@ class TextModel(nn.Module):
 
 class CLIP(TextModel):
     """
-    OpenAI CLIP text encoder implementation.
+    Implements OpenAI's CLIP (Contrastive Language-Image Pre-training) text encoder.
 
-    This class implements the TextModel interface using OpenAI's CLIP model for text encoding.
+    This class provides a text encoder based on OpenAI's CLIP model, which can convert text into feature vectors
+    that are aligned with corresponding image features in a shared embedding space.
 
     Attributes:
         model (clip.model.CLIP): The loaded CLIP model.
@@ -70,15 +71,33 @@ class CLIP(TextModel):
     Methods:
         tokenize: Convert input texts to CLIP tokens.
         encode_text: Encode tokenized texts into normalized feature vectors.
+
+    Examples:
+        >>> from ultralytics.models.sam import CLIP
+        >>> import torch
+        >>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        >>> clip_model = CLIP(size="ViT-B/32", device=device)
+        >>> tokens = clip_model.tokenize(["a photo of a cat", "a photo of a dog"])
+        >>> text_features = clip_model.encode_text(tokens)
+        >>> print(text_features.shape)
     """
 
     def __init__(self, size, device):
         """
         Initialize the CLIP text encoder.
 
+        This class implements the TextModel interface using OpenAI's CLIP model for text encoding. It loads
+        a pre-trained CLIP model of the specified size and prepares it for text encoding tasks.
+
         Args:
             size (str): Model size identifier (e.g., 'ViT-B/32').
             device (torch.device): Device to load the model on.
+
+        Examples:
+            >>> import torch
+            >>> from ultralytics.models.sam.modules.clip import CLIP
+            >>> clip_model = CLIP("ViT-B/32", device=torch.device("cuda:0"))
+            >>> text_features = clip_model.encode_text(["a photo of a cat", "a photo of a dog"])
         """
         super().__init__()
         self.model = clip.load(size, device=device)[0]
@@ -87,7 +106,20 @@ class CLIP(TextModel):
         self.eval()
 
     def tokenize(self, texts):
-        """Convert input texts to CLIP tokens."""
+        """
+        Convert input texts to CLIP tokens.
+
+        Args:
+            texts (str | List[str]): Input text or list of texts to tokenize.
+
+        Returns:
+            (torch.Tensor): Tokenized text tensor with shape (batch_size, context_length) ready for model processing.
+
+        Examples:
+            >>> model = CLIP("ViT-B/32", device="cpu")
+            >>> tokens = model.tokenize("a photo of a cat")
+            >>> print(tokens.shape)  # torch.Size([1, 77])
+        """
         return clip.tokenize(texts).to(self.device)
 
     @smart_inference_mode()
@@ -95,12 +127,22 @@ class CLIP(TextModel):
         """
         Encode tokenized texts into normalized feature vectors.
 
+        This method processes tokenized text inputs through the CLIP model to generate feature vectors, which are then
+        normalized to unit length. These normalized vectors can be used for text-image similarity comparisons.
+
         Args:
-            texts (torch.Tensor): Tokenized text inputs.
-            dtype (torch.dtype): Data type for output features.
+            texts (torch.Tensor): Tokenized text inputs, typically created using the tokenize() method.
+            dtype (torch.dtype, optional): Data type for output features. Default is torch.float32.
 
         Returns:
-            (torch.Tensor): Normalized text feature vectors.
+            (torch.Tensor): Normalized text feature vectors with unit length (L2 norm = 1).
+
+        Examples:
+            >>> clip_model = CLIP("ViT-B/32", device="cuda")
+            >>> tokens = clip_model.tokenize(["a photo of a cat", "a photo of a dog"])
+            >>> features = clip_model.encode_text(tokens)
+            >>> features.shape
+            torch.Size([2, 512])
         """
         txt_feats = self.model.encode_text(texts).to(dtype)
         txt_feats = txt_feats / txt_feats.norm(p=2, dim=-1, keepdim=True)
@@ -109,9 +151,10 @@ class CLIP(TextModel):
 
 class MobileCLIP(TextModel):
     """
-    Apple MobileCLIP text encoder implementation.
+    Implement Apple's MobileCLIP text encoder for efficient text encoding.
 
-    This class implements the TextModel interface using Apple's MobileCLIP model for efficient text encoding.
+    This class implements the TextModel interface using Apple's MobileCLIP model, providing efficient text encoding
+    capabilities for vision-language tasks.
 
     Attributes:
         model (mobileclip.model.MobileCLIP): The loaded MobileCLIP model.
@@ -122,6 +165,12 @@ class MobileCLIP(TextModel):
     Methods:
         tokenize: Convert input texts to MobileCLIP tokens.
         encode_text: Encode tokenized texts into normalized feature vectors.
+
+    Examples:
+        >>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        >>> text_encoder = MobileCLIP(size="s0", device=device)
+        >>> tokens = text_encoder.tokenize(["a photo of a cat", "a photo of a dog"])
+        >>> features = text_encoder.encode_text(tokens)
     """
 
     config_size_map = {"s0": "s0", "s1": "s1", "s2": "s2", "b": "b", "blt": "b"}
@@ -130,9 +179,18 @@ class MobileCLIP(TextModel):
         """
         Initialize the MobileCLIP text encoder.
 
+        This class implements the TextModel interface using Apple's MobileCLIP model for efficient text encoding.
+
         Args:
             size (str): Model size identifier (e.g., 's0', 's1', 's2', 'b', 'blt').
             device (torch.device): Device to load the model on.
+
+        Examples:
+            >>> from ultralytics.nn.modules import MobileCLIP
+            >>> import torch
+            >>> model = MobileCLIP("s0", device=torch.device("cpu"))
+            >>> tokens = model.tokenize(["a photo of a cat", "a photo of a dog"])
+            >>> features = model.encode_text(tokens)
         """
         super().__init__()
         config = self.config_size_map[size]
@@ -148,7 +206,19 @@ class MobileCLIP(TextModel):
         self.eval()
 
     def tokenize(self, texts):
-        """Convert input texts to MobileCLIP tokens."""
+        """
+        Convert input texts to MobileCLIP tokens.
+
+        Args:
+            texts (list[str]): List of text strings to tokenize.
+
+        Returns:
+            (torch.Tensor): Tokenized text inputs with shape (batch_size, sequence_length).
+
+        Examples:
+            >>> model = MobileCLIP("s0", "cpu")
+            >>> tokens = model.tokenize(["a photo of a cat", "a photo of a dog"])
+        """
         return self.tokenizer(texts).to(self.device)
 
     @smart_inference_mode()
@@ -158,10 +228,17 @@ class MobileCLIP(TextModel):
 
         Args:
             texts (torch.Tensor): Tokenized text inputs.
-            dtype (torch.dtype): Data type for output features.
+            dtype (torch.dtype, optional): Data type for output features.
 
         Returns:
-            (torch.Tensor): Normalized text feature vectors.
+            (torch.Tensor): Normalized text feature vectors with L2 normalization applied.
+
+        Examples:
+            >>> model = MobileCLIP("s0", device="cpu")
+            >>> tokens = model.tokenize(["a photo of a cat", "a photo of a dog"])
+            >>> features = model.encode_text(tokens)
+            >>> features.shape
+            torch.Size([2, 512])  # Actual dimension depends on model size
         """
         text_features = self.model.encode_text(texts).to(dtype)
         text_features /= text_features.norm(p=2, dim=-1, keepdim=True)
@@ -179,15 +256,14 @@ def build_text_model(variant, device=None):
     Returns:
         (TextModel): Instantiated text encoding model.
 
-    Raises:
-        AssertionError: If the specified variant is not supported.
+    Examples:
+        >>> model = build_text_model("clip:ViT-B/32", device=torch.device("cuda"))
+        >>> model = build_text_model("mobileclip:s0", device=torch.device("cpu"))
     """
-    LOGGER.info(f"Build text model {variant}")
     base, size = variant.split(":")
     if base == "clip":
         return CLIP(size, device)
     elif base == "mobileclip":
         return MobileCLIP(size, device)
     else:
-        print("Variant not found")
-        assert False
+        raise ValueError(f"Unrecognized base model: '{base}'. Supported base models: 'clip', 'mobileclip'.")

ultralytics/utils/benchmarks.py

@@ -126,7 +126,7 @@ def benchmark(
                 assert not isinstance(model, YOLOWorld), "YOLOWorldv2 TensorFlow exports not supported by onnx2tf yet"
             if i == 11:  # Paddle
                 assert not isinstance(model, YOLOWorld), "YOLOWorldv2 Paddle exports not supported yet"
-                assert not model.task == "obb", "Paddle OBB bug https://github.com/PaddlePaddle/Paddle/issues/72024"
+                assert model.task != "obb", "Paddle OBB bug https://github.com/PaddlePaddle/Paddle/issues/72024"
                 assert not is_end2end, "End-to-end models not supported by PaddlePaddle yet"
                 assert LINUX or MACOS, "Windows Paddle exports not supported yet"
             if i == 12:  # MNN