dgenerate-ultralytics-headless 8.3.194__py3-none-any.whl → 8.3.196__py3-none-any.whl

ultralytics/utils/metrics.py

@@ -318,7 +318,7 @@ class ConfusionMatrix(DataExportMixin):
         task (str): The type of task, either 'detect' or 'classify'.
         matrix (np.ndarray): The confusion matrix, with dimensions depending on the task.
         nc (int): The number of category.
-        names (List[str]): The names of the classes, used as labels on the plot.
+        names (list[str]): The names of the classes, used as labels on the plot.
         matches (dict): Contains the indices of ground truths and predictions categorized into TP, FP and FN.
     """
 
@@ -327,7 +327,7 @@ class ConfusionMatrix(DataExportMixin):
         Initialize a ConfusionMatrix instance.
 
         Args:
-            names (Dict[int, str], optional): Names of classes, used as labels on the plot.
+            names (dict[int, str], optional): Names of classes, used as labels on the plot.
             task (str, optional): Type of task, either 'detect' or 'classify'.
             save_matches (bool, optional): Save the indices of GTs, TPs, FPs, FNs for visualization.
         """
@@ -346,7 +346,7 @@ class ConfusionMatrix(DataExportMixin):
 
         Args:
             mtype (str): Match type identifier ('TP', 'FP', 'FN' or 'GT').
-            batch (Dict[str, Any]): Batch data containing detection results with keys
+            batch (dict[str, Any]): Batch data containing detection results with keys
                 like 'bboxes', 'cls', 'conf', 'keypoints', 'masks'.
             idx (int): Index of the specific detection to append from the batch.
 
@@ -368,8 +368,8 @@ class ConfusionMatrix(DataExportMixin):
         Update confusion matrix for classification task.
 
         Args:
-            preds (List[N, min(nc,5)]): Predicted class labels.
-            targets (List[N, 1]): Ground truth class labels.
+            preds (list[N, min(nc,5)]): Predicted class labels.
+            targets (list[N, 1]): Ground truth class labels.
         """
         preds, targets = torch.cat(preds)[:, 0], torch.cat(targets)
         for p, t in zip(preds.cpu().numpy(), targets.cpu().numpy()):
@@ -386,10 +386,10 @@ class ConfusionMatrix(DataExportMixin):
         Update confusion matrix for object detection task.
 
         Args:
-            detections (Dict[str, torch.Tensor]): Dictionary containing detected bounding boxes and their associated information.
+            detections (dict[str, torch.Tensor]): Dictionary containing detected bounding boxes and their associated information.
                                        Should contain 'cls', 'conf', and 'bboxes' keys, where 'bboxes' can be
                                        Array[N, 4] for regular boxes or Array[N, 5] for OBB with angle.
-            batch (Dict[str, Any]): Batch dictionary containing ground truth data with 'bboxes' (Array[M, 4]| Array[M, 5]) and
+            batch (dict[str, Any]): Batch dictionary containing ground truth data with 'bboxes' (Array[M, 4]| Array[M, 5]) and
                 'cls' (Array[M]) keys, where M is the number of ground truth objects.
             conf (float, optional): Confidence threshold for detections.
             iou_thres (float, optional): IoU threshold for matching detections to ground truth.
@@ -599,7 +599,7 @@ class ConfusionMatrix(DataExportMixin):
             decimals (int): Number of decimal places to round the output values to.
 
         Returns:
-            (List[Dict[str, float]]): A list of dictionaries, each representing one predicted class with corresponding values for all actual classes.
+            (list[dict[str, float]]): A list of dictionaries, each representing one predicted class with corresponding values for all actual classes.
 
         Examples:
             >>> results = model.val(data="coco8.yaml", plots=True)
@@ -651,7 +651,7 @@ def plot_pr_curve(
         py (np.ndarray): Y values for the PR curve.
         ap (np.ndarray): Average precision values.
         save_dir (Path, optional): Path to save the plot.
-        names (Dict[int, str], optional): Dictionary mapping class indices to class names.
+        names (dict[int, str], optional): Dictionary mapping class indices to class names.
         on_plot (callable, optional): Function to call after plot is saved.
     """
     import matplotlib.pyplot as plt  # scope for faster 'import ultralytics'
@@ -695,7 +695,7 @@ def plot_mc_curve(
         px (np.ndarray): X values for the metric-confidence curve.
         py (np.ndarray): Y values for the metric-confidence curve.
         save_dir (Path, optional): Path to save the plot.
-        names (Dict[int, str], optional): Dictionary mapping class indices to class names.
+        names (dict[int, str], optional): Dictionary mapping class indices to class names.
         xlabel (str, optional): X-axis label.
         ylabel (str, optional): Y-axis label.
         on_plot (callable, optional): Function to call after plot is saved.
@@ -780,7 +780,7 @@ def ap_per_class(
         plot (bool, optional): Whether to plot PR curves or not.
         on_plot (callable, optional): A callback to pass plots path and data when they are rendered.
         save_dir (Path, optional): Directory to save the PR curves.
-        names (Dict[int, str], optional): Dictionary of class names to plot PR curves.
+        names (dict[int, str], optional): Dictionary of class names to plot PR curves.
         eps (float, optional): A small value to avoid division by zero.
         prefix (str, optional): A prefix string for saving the plot files.
 
@@ -1034,11 +1034,11 @@ class DetMetrics(SimpleClass, DataExportMixin):
     Utility class for computing detection metrics such as precision, recall, and mean average precision (mAP).
 
     Attributes:
-        names (Dict[int, str]): A dictionary of class names.
+        names (dict[int, str]): A dictionary of class names.
         box (Metric): An instance of the Metric class for storing detection results.
-        speed (Dict[str, float]): A dictionary for storing execution times of different parts of the detection process.
+        speed (dict[str, float]): A dictionary for storing execution times of different parts of the detection process.
         task (str): The task type, set to 'detect'.
-        stats (Dict[str, List]): A dictionary containing lists for true positives, confidence scores, predicted classes, target classes, and target images.
+        stats (dict[str, list]): A dictionary containing lists for true positives, confidence scores, predicted classes, target classes, and target images.
         nt_per_class: Number of targets per class.
         nt_per_image: Number of targets per image.
 
@@ -1063,7 +1063,7 @@ class DetMetrics(SimpleClass, DataExportMixin):
         Initialize a DetMetrics instance with a save directory, plot flag, and class names.
 
         Args:
-            names (Dict[int, str], optional): Dictionary of class names.
+            names (dict[int, str], optional): Dictionary of class names.
         """
         self.names = names
         self.box = Metric()
@@ -1078,7 +1078,7 @@ class DetMetrics(SimpleClass, DataExportMixin):
         Update statistics by appending new values to existing stat collections.
 
         Args:
-            stat (Dict[str, any]): Dictionary containing new statistical values to append.
+            stat (dict[str, any]): Dictionary containing new statistical values to append.
                          Keys should match existing keys in self.stats.
         """
         for k in self.stats.keys():
@@ -1094,7 +1094,7 @@ class DetMetrics(SimpleClass, DataExportMixin):
             on_plot (callable, optional): Function to call after plots are generated. Defaults to None.
 
         Returns:
-            (Dict[str, np.ndarray]): Dictionary containing concatenated statistics arrays.
+            (dict[str, np.ndarray]): Dictionary containing concatenated statistics arrays.
         """
         stats = {k: np.concatenate(v, 0) for k, v in self.stats.items()}  # to numpy
         if not stats:
@@ -1176,7 +1176,7 @@ class DetMetrics(SimpleClass, DataExportMixin):
            decimals (int): Number of decimal places to round the metrics values to.
 
         Returns:
-           (List[Dict[str, Any]]): A list of dictionaries, each representing one class with corresponding metric values.
+           (list[dict[str, Any]]): A list of dictionaries, each representing one class with corresponding metric values.
 
         Examples:
            >>> results = model.val(data="coco8.yaml")
@@ -1206,12 +1206,12 @@ class SegmentMetrics(DetMetrics):
     Calculate and aggregate detection and segmentation metrics over a given set of classes.
 
     Attributes:
-        names (Dict[int, str]): Dictionary of class names.
+        names (dict[int, str]): Dictionary of class names.
         box (Metric): An instance of the Metric class for storing detection results.
         seg (Metric): An instance of the Metric class to calculate mask segmentation metrics.
-        speed (Dict[str, float]): A dictionary for storing execution times of different parts of the detection process.
+        speed (dict[str, float]): A dictionary for storing execution times of different parts of the detection process.
         task (str): The task type, set to 'segment'.
-        stats (Dict[str, List]): A dictionary containing lists for true positives, confidence scores, predicted classes, target classes, and target images.
+        stats (dict[str, list]): A dictionary containing lists for true positives, confidence scores, predicted classes, target classes, and target images.
         nt_per_class: Number of targets per class.
         nt_per_image: Number of targets per image.
 
@@ -1232,7 +1232,7 @@ class SegmentMetrics(DetMetrics):
         Initialize a SegmentMetrics instance with a save directory, plot flag, and class names.
 
         Args:
-            names (Dict[int, str], optional): Dictionary of class names.
+            names (dict[int, str], optional): Dictionary of class names.
         """
         DetMetrics.__init__(self, names)
         self.seg = Metric()
@@ -1249,7 +1249,7 @@ class SegmentMetrics(DetMetrics):
             on_plot (callable, optional): Function to call after plots are generated. Defaults to None.
 
         Returns:
-            (Dict[str, np.ndarray]): Dictionary containing concatenated statistics arrays.
+            (dict[str, np.ndarray]): Dictionary containing concatenated statistics arrays.
         """
         stats = DetMetrics.process(self, save_dir, plot, on_plot=on_plot)  # process box stats
         results_mask = ap_per_class(
@@ -1320,7 +1320,7 @@ class SegmentMetrics(DetMetrics):
             decimals (int): Number of decimal places to round the metrics values to.
 
         Returns:
-            (List[Dict[str, Any]]): A list of dictionaries, each representing one class with corresponding metric values.
+            (list[dict[str, Any]]): A list of dictionaries, each representing one class with corresponding metric values.
 
         Examples:
             >>> results = model.val(data="coco8-seg.yaml")
@@ -1343,12 +1343,12 @@ class PoseMetrics(DetMetrics):
     Calculate and aggregate detection and pose metrics over a given set of classes.
 
     Attributes:
-        names (Dict[int, str]): Dictionary of class names.
+        names (dict[int, str]): Dictionary of class names.
         pose (Metric): An instance of the Metric class to calculate pose metrics.
         box (Metric): An instance of the Metric class for storing detection results.
-        speed (Dict[str, float]): A dictionary for storing execution times of different parts of the detection process.
+        speed (dict[str, float]): A dictionary for storing execution times of different parts of the detection process.
         task (str): The task type, set to 'pose'.
-        stats (Dict[str, List]): A dictionary containing lists for true positives, confidence scores, predicted classes, target classes, and target images.
+        stats (dict[str, list]): A dictionary containing lists for true positives, confidence scores, predicted classes, target classes, and target images.
         nt_per_class: Number of targets per class.
         nt_per_image: Number of targets per image.
 
@@ -1369,7 +1369,7 @@ class PoseMetrics(DetMetrics):
         Initialize the PoseMetrics class with directory path, class names, and plotting options.
 
         Args:
-            names (Dict[int, str], optional): Dictionary of class names.
+            names (dict[int, str], optional): Dictionary of class names.
         """
         super().__init__(names)
         self.pose = Metric()
@@ -1386,7 +1386,7 @@ class PoseMetrics(DetMetrics):
             on_plot (callable, optional): Function to call after plots are generated.
 
         Returns:
-            (Dict[str, np.ndarray]): Dictionary containing concatenated statistics arrays.
+            (dict[str, np.ndarray]): Dictionary containing concatenated statistics arrays.
         """
         stats = DetMetrics.process(self, save_dir, plot, on_plot=on_plot)  # process box stats
         results_pose = ap_per_class(
@@ -1461,7 +1461,7 @@ class PoseMetrics(DetMetrics):
             decimals (int): Number of decimal places to round the metrics values to.
 
         Returns:
-            (List[Dict[str, Any]]): A list of dictionaries, each representing one class with corresponding metric values.
+            (list[dict[str, Any]]): A list of dictionaries, each representing one class with corresponding metric values.
 
         Examples:
             >>> results = model.val(data="coco8-pose.yaml")
@@ -1553,7 +1553,7 @@ class ClassifyMetrics(SimpleClass, DataExportMixin):
             decimals (int): Number of decimal places to round the metrics values to.
 
         Returns:
-            (List[Dict[str, float]]): A list with one dictionary containing Top-1 and Top-5 classification accuracy.
+            (list[dict[str, float]]): A list with one dictionary containing Top-1 and Top-5 classification accuracy.
 
         Examples:
             >>> results = model.val(data="imagenet10")
@@ -1568,11 +1568,11 @@ class OBBMetrics(DetMetrics):
     Metrics for evaluating oriented bounding box (OBB) detection.
 
     Attributes:
-        names (Dict[int, str]): Dictionary of class names.
+        names (dict[int, str]): Dictionary of class names.
         box (Metric): An instance of the Metric class for storing detection results.
-        speed (Dict[str, float]): A dictionary for storing execution times of different parts of the detection process.
+        speed (dict[str, float]): A dictionary for storing execution times of different parts of the detection process.
         task (str): The task type, set to 'obb'.
-        stats (Dict[str, List]): A dictionary containing lists for true positives, confidence scores, predicted classes, target classes, and target images.
+        stats (dict[str, list]): A dictionary containing lists for true positives, confidence scores, predicted classes, target classes, and target images.
         nt_per_class: Number of targets per class.
         nt_per_image: Number of targets per image.
 
@@ -1585,7 +1585,7 @@ class OBBMetrics(DetMetrics):
         Initialize an OBBMetrics instance with directory, plotting, and class names.
 
         Args:
-            names (Dict[int, str], optional): Dictionary of class names.
+            names (dict[int, str], optional): Dictionary of class names.
         """
         DetMetrics.__init__(self, names)
         # TODO: probably remove task as well

ultralytics/utils/nms.py

@@ -38,10 +38,10 @@ def non_max_suppression(
             containing boxes, classes, and optional masks.
         conf_thres (float): Confidence threshold for filtering detections. Valid values are between 0.0 and 1.0.
         iou_thres (float): IoU threshold for NMS filtering. Valid values are between 0.0 and 1.0.
-        classes (List[int], optional): List of class indices to consider. If None, all classes are considered.
+        classes (list[int], optional): List of class indices to consider. If None, all classes are considered.
         agnostic (bool): Whether to perform class-agnostic NMS.
         multi_label (bool): Whether each box can have multiple labels.
-        labels (List[List[Union[int, float, torch.Tensor]]]): A priori labels for each image.
+        labels (list[list[Union[int, float, torch.Tensor]]]): A priori labels for each image.
         max_det (int): Maximum number of detections to keep per image.
         nc (int): Number of classes. Indices after this are considered masks.
         max_time_img (float): Maximum time in seconds for processing one image.
@@ -52,9 +52,9 @@ def non_max_suppression(
         return_idxs (bool): Whether to return the indices of kept detections.
 
     Returns:
-        output (List[torch.Tensor]): List of detections per image with shape (num_boxes, 6 + num_masks)
+        output (list[torch.Tensor]): List of detections per image with shape (num_boxes, 6 + num_masks)
             containing (x1, y1, x2, y2, confidence, class, mask1, mask2, ...).
-        keepi (List[torch.Tensor]): Indices of kept detections if return_idxs=True.
+        keepi (list[torch.Tensor]): Indices of kept detections if return_idxs=True.
     """
     # Checks
     assert 0 <= conf_thres <= 1, f"Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0"

ultralytics/utils/ops.py

@@ -135,7 +135,7 @@ def scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None, padding: bool = T
             boxes[..., 2] -= pad_x  # x padding
             boxes[..., 3] -= pad_y  # y padding
     boxes[..., :4] /= gain
-    return clip_boxes(boxes, img0_shape)
+    return boxes if xywh else clip_boxes(boxes, img0_shape)
 
 
 def make_divisible(x: int, divisor):

ultralytics/utils/patches.py

@@ -51,7 +51,7 @@ def imwrite(filename: str, img: np.ndarray, params: list[int] | None = None) ->
     Args:
         filename (str): Path to the file to write.
         img (np.ndarray): Image to write.
-        params (List[int], optional): Additional parameters for image encoding.
+        params (list[int], optional): Additional parameters for image encoding.
 
     Returns:
         (bool): True if the file was written successfully, False otherwise.
@@ -172,7 +172,7 @@ def override_configs(args, overrides: dict[str, Any] | None = None):
 
     Args:
         args (IterableSimpleNamespace): Original configuration arguments.
-        overrides (Dict[str, Any]): Dictionary of overrides to apply.
+        overrides (dict[str, Any]): Dictionary of overrides to apply.
 
     Yields:
         (IterableSimpleNamespace): Configuration arguments with overrides applied.

ultralytics/utils/plotting.py

@@ -26,7 +26,7 @@ class Colors:
     RGB values and accessing predefined color schemes for object detection and pose estimation.
 
     Attributes:
-        palette (List[tuple]): List of RGB color tuples for general use.
+        palette (list[tuple]): List of RGB color tuples for general use.
         n (int): The number of colors in the palette.
         pose_palette (np.ndarray): A specific color palette array for pose estimation with dtype np.uint8.
 
@@ -176,9 +176,9 @@ class Annotator:
         pil (bool): Whether to use PIL or cv2 for drawing annotations.
         font (ImageFont.truetype | ImageFont.load_default): Font used for text annotations.
         lw (float): Line width for drawing.
-        skeleton (List[List[int]]): Skeleton structure for keypoints.
-        limb_color (List[int]): Color palette for limbs.
-        kpt_color (List[int]): Color palette for keypoints.
+        skeleton (list[list[int]]): Skeleton structure for keypoints.
+        limb_color (list[int]): Color palette for limbs.
+        kpt_color (list[int]): Color palette for keypoints.
         dark_colors (set): Set of colors considered dark for text contrast.
         light_colors (set): Set of colors considered light for text contrast.
 
@@ -369,7 +369,7 @@ class Annotator:
 
         Args:
             masks (torch.Tensor | np.ndarray): Predicted masks with shape: [n, h, w]
-            colors (List[List[int]]): Colors for predicted masks, [[r, g, b] * n]
+            colors (list[list[int]]): Colors for predicted masks, [[r, g, b] * n]
             im_gpu (torch.Tensor | None): Image is in cuda, shape: [3, h, w], range: [0, 1]
             alpha (float, optional): Mask transparency: 0.0 fully transparent, 1.0 opaque.
             retina_masks (bool, optional): Whether to use high resolution masks or not.
@@ -484,7 +484,7 @@ class Annotator:
         Add text to an image using PIL or cv2.
 
         Args:
-            xy (List[int]): Top-left coordinates for text placement.
+            xy (list[int]): Top-left coordinates for text placement.
             text (str): Text to be drawn.
             txt_color (tuple, optional): Text color (R, G, B).
             anchor (str, optional): Text anchor position ('top' or 'bottom').
@@ -695,11 +695,11 @@ def plot_images(
     Plot image grid with labels, bounding boxes, masks, and keypoints.
 
     Args:
-        labels (Dict[str, Any]): Dictionary containing detection data with keys like 'cls', 'bboxes', 'conf', 'masks', 'keypoints', 'batch_idx', 'img'.
+        labels (dict[str, Any]): Dictionary containing detection data with keys like 'cls', 'bboxes', 'conf', 'masks', 'keypoints', 'batch_idx', 'img'.
         images (torch.Tensor | np.ndarray]): Batch of images to plot. Shape: (batch_size, channels, height, width).
-        paths (Optional[List[str]]): List of file paths for each image in the batch.
+        paths (Optional[list[str]]): List of file paths for each image in the batch.
         fname (str): Output filename for the plotted image grid.
-        names (Optional[Dict[int, str]]): Dictionary mapping class indices to class names.
+        names (Optional[dict[int, str]]): Dictionary mapping class indices to class names.
         on_plot (Optional[Callable]): Optional callback function to be called after saving the plot.
         max_size (int): Maximum size of the output image grid.
         max_subplots (int): Maximum number of subplots in the image grid.
@@ -1004,6 +1004,7 @@ def plot_tune_results(csv_file: str = "tune_results.csv"):
     _save_one_file(csv_file.with_name("tune_fitness.png"))
 
 
+@plt_settings()
 def feature_visualization(x, module_type: str, stage: int, n: int = 32, save_dir: Path = Path("runs/detect/exp")):
     """
     Visualize feature maps of a given model module during inference.

ultralytics/utils/torch_utils.py

@@ -33,6 +33,7 @@ from ultralytics.utils import (
     colorstr,
 )
 from ultralytics.utils.checks import check_version
+from ultralytics.utils.cpu import CPUInfo
 from ultralytics.utils.patches import torch_load
 
 # Version checks (all default to version>=min_version)
@@ -113,12 +114,7 @@ def get_cpu_info():
 
     if "cpu_info" not in PERSISTENT_CACHE:
         try:
-            import cpuinfo  # pip install py-cpuinfo
-
-            k = "brand_raw", "hardware_raw", "arch_string_raw"  # keys sorted by preference
-            info = cpuinfo.get_cpu_info()  # info dict
-            string = info.get(k[0] if k[0] in info else k[1] if k[1] in info else k[2], "unknown")
-            PERSISTENT_CACHE["cpu_info"] = string.replace("(R)", "").replace("CPU ", "").replace("@ ", "")
+            PERSISTENT_CACHE["cpu_info"] = CPUInfo.name()
         except Exception:
             pass
     return PERSISTENT_CACHE.get("cpu_info", "unknown")
@@ -433,7 +429,7 @@ def get_flops(model, imgsz=640):
         return 0.0  # if not installed return 0.0 GFLOPs
 
     try:
-        model = de_parallel(model)
+        model = unwrap_model(model)
         p = next(model.parameters())
         if not isinstance(imgsz, list):
             imgsz = [imgsz, imgsz]  # expand if int/float
@@ -464,7 +460,7 @@ def get_flops_with_torch_profiler(model, imgsz=640):
     """
     if not TORCH_2_0:  # torch profiler implemented in torch>=2.0
         return 0.0
-    model = de_parallel(model)
+    model = unwrap_model(model)
     p = next(model.parameters())
     if not isinstance(imgsz, list):
         imgsz = [imgsz, imgsz]  # expand if int/float
@@ -581,17 +577,24 @@ def is_parallel(model):
     return isinstance(model, (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel))
 
 
-def de_parallel(model):
+def unwrap_model(m: nn.Module) -> nn.Module:
     """
-    De-parallelize a model: return single-GPU model if model is of type DP or DDP.
+    Unwrap compiled and parallel models to get the base model.
 
     Args:
-        model (nn.Module): Model to de-parallelize.
+        m (nn.Module): A model that may be wrapped by torch.compile (._orig_mod) or parallel wrappers such as
+            DataParallel/DistributedDataParallel (.module).
 
     Returns:
-        (nn.Module): De-parallelized model.
+        m (nn.Module): The unwrapped base model without compile or parallel wrappers.
     """
-    return model.module if is_parallel(model) else model
+    while True:
+        if hasattr(m, "_orig_mod") and isinstance(m._orig_mod, nn.Module):
+            m = m._orig_mod
+        elif hasattr(m, "module") and isinstance(m.module, nn.Module):
+            m = m.module
+        else:
+            return m
 
 
 def one_cycle(y1=0.0, y2=1.0, steps=100):
@@ -673,7 +676,7 @@ class ModelEMA:
             tau (int, optional): EMA decay time constant.
             updates (int, optional): Initial number of updates.
         """
-        self.ema = deepcopy(de_parallel(model)).eval()  # FP32 EMA
+        self.ema = deepcopy(unwrap_model(model)).eval()  # FP32 EMA
         self.updates = updates  # number of EMA updates
         self.decay = lambda x: decay * (1 - math.exp(-x / tau))  # decay exponential ramp (to help early epochs)
         for p in self.ema.parameters():
@@ -691,7 +694,7 @@ class ModelEMA:
             self.updates += 1
             d = self.decay(self.updates)
 
-            msd = de_parallel(model).state_dict()  # model state_dict
+            msd = unwrap_model(model).state_dict()  # model state_dict
             for k, v in self.ema.state_dict().items():
                 if v.dtype.is_floating_point:  # true for FP16 and FP32
                     v *= d
@@ -1001,3 +1004,98 @@ class FXModel(nn.Module):
             x = m(x)  # run
             y.append(x)  # save output
         return x
+
+
+def disable_dynamo(func: Any) -> Any:
+    """
+    Disable torch.compile/dynamo for a callable when available.
+
+    Args:
+        func (Any): Callable object to wrap. Could be a function, method, or class.
+
+    Returns:
+        func (Any): Same callable, wrapped by torch._dynamo.disable when available, otherwise unchanged.
+
+    Examples:
+        >>> @disable_dynamo
+        ... def fn(x):
+        ...     return x + 1
+        >>> # Works even if torch._dynamo is not available
+        >>> _ = fn(1)
+    """
+    if hasattr(torch, "_dynamo"):
+        return torch._dynamo.disable(func)
+    return func
+
+
+def attempt_compile(
+    model: torch.nn.Module,
+    device: torch.device,
+    imgsz: int = 640,
+    use_autocast: bool = False,
+    warmup: bool = False,
+    prefix: str = colorstr("compile:"),
+) -> torch.nn.Module:
+    """
+    Compile a model with torch.compile and optionally warm up the graph to reduce first-iteration latency.
+
+    This utility attempts to compile the provided model using the inductor backend with dynamic shapes enabled and an
+    autotuning mode. If compilation is unavailable or fails, the original model is returned unchanged. An optional
+    warmup performs a single forward pass on a dummy input to prime the compiled graph and measure compile/warmup time.
+
+    Args:
+        model (torch.nn.Module): Model to compile.
+        device (torch.device): Inference device used for warmup and autocast decisions.
+        imgsz (int, optional): Square input size to create a dummy tensor with shape (1, 3, imgsz, imgsz) for warmup.
+        use_autocast (bool, optional): Whether to run warmup under autocast on CUDA or MPS devices.
+        warmup (bool, optional): Whether to execute a single dummy forward pass to warm up the compiled model.
+        prefix (str, optional): Message prefix for logger output.
+
+    Returns:
+        model (torch.nn.Module): Compiled model if compilation succeeds, otherwise the original unmodified model.
+
+    Notes:
+        - If the current PyTorch build does not provide torch.compile, the function returns the input model immediately.
+        - Warmup runs under torch.inference_mode and may use torch.autocast for CUDA/MPS to align compute precision.
+        - CUDA devices are synchronized after warmup to account for asynchronous kernel execution.
+
+    Examples:
+        >>> device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        >>> # Try to compile and warm up a model with a 640x640 input
+        >>> model = attempt_compile(model, device=device, imgsz=640, use_autocast=True, warmup=True)
+    """
+    if not hasattr(torch, "compile"):
+        return model
+
+    LOGGER.info(f"{prefix} starting torch.compile...")
+    t0 = time.perf_counter()
+    try:
+        model = torch.compile(model, mode="max-autotune", backend="inductor")
+    except Exception as e:
+        LOGGER.warning(f"{prefix} torch.compile failed, continuing uncompiled: {e}")
+        return model
+    t_compile = time.perf_counter() - t0
+
+    t_warm = 0.0
+    if warmup:
+        # Use a single dummy tensor to build the graph shape state and reduce first-iteration latency
+        dummy = torch.zeros(1, 3, imgsz, imgsz, device=device)
+        if use_autocast and device.type == "cuda":
+            dummy = dummy.half()
+        t1 = time.perf_counter()
+        with torch.inference_mode():
+            if use_autocast and device.type in {"cuda", "mps"}:
+                with torch.autocast(device.type):
+                    _ = model(dummy)
+            else:
+                _ = model(dummy)
+        if device.type == "cuda":
+            torch.cuda.synchronize(device)
+        t_warm = time.perf_counter() - t1
+
+    total = t_compile + t_warm
+    if warmup:
+        LOGGER.info(f"{prefix} complete in {total:.1f}s (compile {t_compile:.1f}s + warmup {t_warm:.1f}s)")
+    else:
+        LOGGER.info(f"{prefix} compile complete in {t_compile:.1f}s (no warmup)")
+    return model

ultralytics/utils/triton.py

@@ -20,10 +20,10 @@ class TritonRemoteModel:
         triton_client: The Triton client (either HTTP or gRPC).
         InferInput: The input class for the Triton client.
         InferRequestedOutput: The output request class for the Triton client.
-        input_formats (List[str]): The data types of the model inputs.
-        np_input_formats (List[type]): The numpy data types of the model inputs.
-        input_names (List[str]): The names of the model inputs.
-        output_names (List[str]): The names of the model outputs.
+        input_formats (list[str]): The data types of the model inputs.
+        np_input_formats (list[type]): The numpy data types of the model inputs.
+        input_names (list[str]): The names of the model inputs.
+        output_names (list[str]): The names of the model outputs.
         metadata: The metadata associated with the model.
 
     Methods:
@@ -96,7 +96,7 @@ class TritonRemoteModel:
                 for the corresponding model input.
 
         Returns:
-            (List[np.ndarray]): Model outputs with the same dtype as the input. Each element in the list
+            (list[np.ndarray]): Model outputs with the same dtype as the input. Each element in the list
                 corresponds to one of the model's output tensors.
 
         Examples: