trt ¶

@classmethod
def autoselect(cls, device_id: int = 0, **kwargs: Any) -> Backend:
    """
    Try to select the best backend for the current system.

    Args:
        device_id: The GPU device id.
        **kwargs: Additional arguments to pass to the backend.

    Returns:
        The selected backend.
    """

    gpu = get_gpu(device_id)
    vendor = (
        cast(str | None, gpu.vendor)
        if gpu
        else "apple"
        # macOS x86_64 is unsupported
        if platform.system().lower() == "darwin" and platform.machine() == "x86_64"
        else None
    )

    match vendor:
        # Windows & Linux
        case "nvidia":
            if hasattr(core, "trt"):
                backend = Backend.TRT
            elif hasattr(core, "trt_rtx"):
                backend = Backend.TRT_RTX
            elif platform.system().lower() == "windows" and hasattr(core, "ort"):
                backend = Backend.ORT_DML
            elif hasattr(core, "ort"):
                backend = Backend.ORT_CUDA
            elif hasattr(core, "ncnn"):
                backend = Backend.NCNN
            else:
                backend = Backend.OV_CPU
        # Windows & Linux
        case "amd":
            if platform.system().lower() == "windows" and hasattr(core, "ort"):
                backend = Backend.ORT_DML
            elif hasattr(core, "migx"):
                backend = Backend.MIGX
            elif hasattr(core, "ncnn"):
                backend = Backend.NCNN_VK
            else:
                backend = Backend.OV_CPU
        # Windows & Linux
        case "intel":
            # device-smi can't detect Intel NPUs in 0.5.6
            # https://github.com/ModelCloud/Device-SMI#roadmap
            if hasattr(core, "ov"):
                backend = Backend.OV_GPU
            elif platform.system().lower() == "windows" and hasattr(core, "ort"):
                backend = Backend.ORT_DML
            elif hasattr(core, "ncnn"):
                backend = Backend.NCNN_VK
            else:
                backend = Backend.OV_CPU
        # macOS ARM64 & x86_64
        case "apple":
            if hasattr(core, "ncnn"):
                backend = Backend.NCNN_VK
            elif hasattr(core, "ort"):
                backend = Backend.ORT_COREML
            else:
                backend = Backend.OV_CPU
        case _:
            backend = Backend.OV_CPU

    del gpu

    return backend(**kwargs)

def build(
    self,
    network_path: Path,
    engine_path: Path,
    channels: int,
    tilesize: Shape,
    input_name: str,
) -> None:
    trt_logger = self.logger
    builder = self.trt.Builder(trt_logger)

    if self.max_threads is not None:
        builder.max_threads = self.max_threads

    network = builder.create_network()
    parser = self.trt.OnnxParser(network, trt_logger)

    if not parser.parse_from_file(str(network_path)):
        errors = [str(parser.get_error(i)) for i in range(parser.num_errors)]
        raise CustomRuntimeError(f"Failed to parse ONNX model: {network_path}\n" + "\n".join(errors))

    config = builder.create_builder_config()

    # Delegate builder setup
    self.configure_builder_config(config, network)
    self.setup_optimization_profile(builder, network, config, channels, input_name, tilesize)

    # Timing Cache
    timing_cache_path = Path(f"{engine_path}.cache")
    timing_cache_data = b""
    if timing_cache_path.exists():
        timing_cache_data = timing_cache_path.read_bytes()

    timing_cache = config.create_timing_cache(timing_cache_data)
    config.set_timing_cache(timing_cache, ignore_mismatch=True)

    # Build
    logger.info(f"Building TensorRT {self.__class__.__name__} engine from {network_path}...")
    serialized = builder.build_serialized_network(network, config)

    if not serialized:
        raise CustomRuntimeError(f"TensorRT engine build failed for {network_path}")

    engine_path.write_bytes(serialized)

    # Save Cache
    updated_cache = config.get_timing_cache()
    timing_cache_path.write_bytes(updated_cache.serialize())

    logger.info(f"Engine saved to {engine_path}")

def build_engine(self, network_path: Path, channels: int, tilesize: Shape, input_name: str = "input") -> Path:
    """
    Build or retrieve a cached TensorRT engine.

    Args:
        network_path: Path to the ONNX model.
        channels: Number of model input channels.
        tilesize: Inference tile size as `(width, height)`.
        input_name: Name of the model input tensor.

    Returns:
        Path to the serialized engine file.
    """
    if self.fp16:
        network_path = self._convert_onnx_fp16(network_path)
    elif self.bf16:
        network_path = self._convert_onnx_bf16(network_path)

    dirname = get_artifacts_folder()
    dirname.mkdir(parents=True, exist_ok=True)
    identity = self.get_identity(network_path, channels, tilesize)
    engine_path = dirname / f"{identity}.engine"

    if not self.force_rebuild and engine_path.is_file() and engine_path.stat().st_size >= 1024:
        return engine_path

    self.build(
        network_path=network_path,
        engine_path=engine_path,
        channels=channels,
        tilesize=tilesize,
        input_name=input_name,
    )

    return engine_path

def configure_builder_config(self, config: trt.IBuilderConfig, network: trt.INetworkDefinition) -> None:
    if self.workspace is not None:
        config.set_memory_pool_limit(self.trt.MemoryPoolType.WORKSPACE, self.workspace)

    if self.tactic_dram is not None:
        config.set_memory_pool_limit(self.trt.MemoryPoolType.TACTIC_DRAM, self.tactic_dram)

    if not self.tf32:
        config.flags &= ~(1 << self.trt.BuilderFlag.TF32.value)

    if self.sparse_weights:
        config.flags |= 1 << self.trt.BuilderFlag.SPARSE_WEIGHTS.value

    if self.strict_nans:
        config.flags |= 1 << self.trt.BuilderFlag.STRICT_NANS.value

    if self.weight_streaming:
        config.flags |= 1 << self.trt.BuilderFlag.WEIGHT_STREAMING.value

    self.configure_tactic_sources(config)
    self.configure_optimization_settings(config)

def configure_optimization_settings(self, config: trt.IBuilderConfig) -> None:
    config.builder_optimization_level = self.builder_optimization_level
    config.avg_timing_iterations = self.avg_timing_iterations

    if self.max_aux_streams is not None:
        config.max_aux_streams = self.max_aux_streams

    if self.max_num_tactics is not None:
        config.max_num_tactics = self.max_num_tactics

    if int(self.tiling_optimization_level) != 0:
        config.tiling_optimization_level = self.trt.TilingOptimizationLevel(self.tiling_optimization_level)
        config.l2_limit_for_tiling = self.l2_limit_for_tiling

def configure_tactic_sources(self, config: trt.IBuilderConfig) -> None:
    tactic_sources = config.get_tactic_sources()

    if self.edge_mask_convolutions:
        tactic_sources |= 1 << self.trt.TacticSource.EDGE_MASK_CONVOLUTIONS.value
    else:
        tactic_sources &= ~(1 << self.trt.TacticSource.EDGE_MASK_CONVOLUTIONS.value)

    if self.jit_convolutions:
        tactic_sources |= 1 << self.trt.TacticSource.JIT_CONVOLUTIONS.value
    else:
        tactic_sources &= ~(1 << self.trt.TacticSource.JIT_CONVOLUTIONS.value)

    config.set_tactic_sources(tactic_sources)

def get_args(self, clips: vs.VideoNode | Sequence[vs.VideoNode]) -> dict[str, Any]:
    return {
        "device_id": self.device_id,
        "use_cuda_graph": self.use_cuda_graph,
        "num_streams": self.num_streams,
        "verbosity": self.verbosity,
    }

def get_identity(self, network_path: Path, channels: int, tilesize: Shape) -> int:
    checksum = zlib.crc32(network_path.read_bytes())

    command = [
        "nvidia-smi",
        "-i",
        str(self.device_id),
        "--query-gpu=name,driver_version",
        "--format=csv,noheader,nounits",
    ]
    res = subprocess.run(command, capture_output=True, text=True, check=True)
    device = [d.strip().replace(" ", "_") for d in res.stdout.split(",")]

    components = (
        str(self),
        str(self.version),
        str(sys.version_info[:2]),
        network_path.name,
        f"{checksum:x}",
        str(channels),
        str(tilesize),
        *device,
    )
    return zlib.crc32(bytes("|".join(components), "utf-8"))

@copy_signature(Backend.inference)
def inference(
    self,
    clips: vs.VideoNode | Sequence[vs.VideoNode],
    network_path: str | os.PathLike[str],
    /,
    overlap: tuple[int, int],
    tilesize: tuple[int, int],
    *,
    flexible: bool = False,
    **kwargs: Any,
) -> vs.VideoNode | list[vs.VideoNode]:
    clips = to_arr(clips)
    channels = sum(clip.format.num_planes for clip in clips)
    engine_path = self.build_engine(Path(network_path), channels, tilesize)

    if self.fp16 or self.bf16:
        # Clips must be in fp16 format is fp16 or bf16 mode is enabled,
        # otherwise the TRT plugins error out.
        clips = [depth(c, 16, sample_type=vs.SampleType.FLOAT) for c in clips]
    else:
        clips = [depth(c, 32) for c in clips]

    return super().inference(clips, engine_path, overlap, tilesize, flexible=flexible, **kwargs)

@classproperty.cached
@classmethod
def logger(cls) -> tensorrt.ILogger:
    from ._trt import Logger

    return Logger(logger)

def setup_optimization_profile(
    self,
    builder: trt.Builder,
    network: trt.INetworkDefinition,
    config: trt.IBuilderConfig,
    channels: int,
    input_name: str,
    tilesize: Shape,
) -> None:
    profile = builder.create_optimization_profile()
    opt_shapes = self.trt.Dims(self.opt_shapes or tilesize)
    max_shapes = self.trt.Dims(self.max_shapes or tilesize)

    input_names = [network.get_input(i).name for i in range(network.num_inputs)]
    if input_name not in input_names:
        logger.debug("input_name %r isn't in the input network", input_name)
        if network.num_inputs == 1:
            input_name = input_names[0]
        else:
            raise CustomValueError(f"Input name '{input_name}' not found in network inputs: {input_names}")

    if self.static_shape:
        shape = self.trt.Dims((1, channels, opt_shapes[1], opt_shapes[0]))

        for i in range(network.num_inputs):
            input_tensor = network.get_input(i)
            if input_tensor.name == input_name:
                input_tensor.shape = shape

        profile.set_shape(input_name, shape, shape, shape)
    else:
        profile.set_shape(
            input_name,
            self.trt.Dims((1, channels, self.min_shapes[1], self.min_shapes[0])),
            self.trt.Dims((1, channels, opt_shapes[1], opt_shapes[0])),
            self.trt.Dims((1, channels, max_shapes[1], max_shapes[0])),
        )

    config.add_optimization_profile(profile)

@classmethod
def autoselect(cls, device_id: int = 0, **kwargs: Any) -> Backend:
    """
    Try to select the best backend for the current system.

    Args:
        device_id: The GPU device id.
        **kwargs: Additional arguments to pass to the backend.

    Returns:
        The selected backend.
    """

    gpu = get_gpu(device_id)
    vendor = (
        cast(str | None, gpu.vendor)
        if gpu
        else "apple"
        # macOS x86_64 is unsupported
        if platform.system().lower() == "darwin" and platform.machine() == "x86_64"
        else None
    )

    match vendor:
        # Windows & Linux
        case "nvidia":
            if hasattr(core, "trt"):
                backend = Backend.TRT
            elif hasattr(core, "trt_rtx"):
                backend = Backend.TRT_RTX
            elif platform.system().lower() == "windows" and hasattr(core, "ort"):
                backend = Backend.ORT_DML
            elif hasattr(core, "ort"):
                backend = Backend.ORT_CUDA
            elif hasattr(core, "ncnn"):
                backend = Backend.NCNN
            else:
                backend = Backend.OV_CPU
        # Windows & Linux
        case "amd":
            if platform.system().lower() == "windows" and hasattr(core, "ort"):
                backend = Backend.ORT_DML
            elif hasattr(core, "migx"):
                backend = Backend.MIGX
            elif hasattr(core, "ncnn"):
                backend = Backend.NCNN_VK
            else:
                backend = Backend.OV_CPU
        # Windows & Linux
        case "intel":
            # device-smi can't detect Intel NPUs in 0.5.6
            # https://github.com/ModelCloud/Device-SMI#roadmap
            if hasattr(core, "ov"):
                backend = Backend.OV_GPU
            elif platform.system().lower() == "windows" and hasattr(core, "ort"):
                backend = Backend.ORT_DML
            elif hasattr(core, "ncnn"):
                backend = Backend.NCNN_VK
            else:
                backend = Backend.OV_CPU
        # macOS ARM64 & x86_64
        case "apple":
            if hasattr(core, "ncnn"):
                backend = Backend.NCNN_VK
            elif hasattr(core, "ort"):
                backend = Backend.ORT_COREML
            else:
                backend = Backend.OV_CPU
        case _:
            backend = Backend.OV_CPU

    del gpu

    return backend(**kwargs)

def build(
    self,
    network_path: Path,
    engine_path: Path,
    channels: int,
    tilesize: Shape,
    input_name: str,
) -> None:
    trt_logger = self.logger
    builder = self.trt.Builder(trt_logger)

    if self.max_threads is not None:
        builder.max_threads = self.max_threads

    network = builder.create_network()
    parser = self.trt.OnnxParser(network, trt_logger)

    if not parser.parse_from_file(str(network_path)):
        errors = [str(parser.get_error(i)) for i in range(parser.num_errors)]
        raise CustomRuntimeError(f"Failed to parse ONNX model: {network_path}\n" + "\n".join(errors))

    config = builder.create_builder_config()

    # Delegate builder setup
    self.configure_builder_config(config, network)
    self.setup_optimization_profile(builder, network, config, channels, input_name, tilesize)

    # Timing Cache
    timing_cache_path = Path(f"{engine_path}.cache")
    timing_cache_data = b""
    if timing_cache_path.exists():
        timing_cache_data = timing_cache_path.read_bytes()

    timing_cache = config.create_timing_cache(timing_cache_data)
    config.set_timing_cache(timing_cache, ignore_mismatch=True)

    # Build
    logger.info(f"Building TensorRT {self.__class__.__name__} engine from {network_path}...")
    serialized = builder.build_serialized_network(network, config)

    if not serialized:
        raise CustomRuntimeError(f"TensorRT engine build failed for {network_path}")

    engine_path.write_bytes(serialized)

    # Save Cache
    updated_cache = config.get_timing_cache()
    timing_cache_path.write_bytes(updated_cache.serialize())

    logger.info(f"Engine saved to {engine_path}")

def build_engine(self, network_path: Path, channels: int, tilesize: Shape, input_name: str = "input") -> Path:
    """
    Build or retrieve a cached TensorRT engine.

    Args:
        network_path: Path to the ONNX model.
        channels: Number of model input channels.
        tilesize: Inference tile size as `(width, height)`.
        input_name: Name of the model input tensor.

    Returns:
        Path to the serialized engine file.
    """
    if self.fp16:
        network_path = self._convert_onnx_fp16(network_path)
    elif self.bf16:
        network_path = self._convert_onnx_bf16(network_path)

    dirname = get_artifacts_folder()
    dirname.mkdir(parents=True, exist_ok=True)
    identity = self.get_identity(network_path, channels, tilesize)
    engine_path = dirname / f"{identity}.engine"

    if not self.force_rebuild and engine_path.is_file() and engine_path.stat().st_size >= 1024:
        return engine_path

    self.build(
        network_path=network_path,
        engine_path=engine_path,
        channels=channels,
        tilesize=tilesize,
        input_name=input_name,
    )

    return engine_path

def configure_builder_config(self, config: trt.IBuilderConfig, network: trt.INetworkDefinition) -> None:
    if self.workspace is not None:
        config.set_memory_pool_limit(self.trt.MemoryPoolType.WORKSPACE, self.workspace)

    if self.tactic_dram is not None:
        config.set_memory_pool_limit(self.trt.MemoryPoolType.TACTIC_DRAM, self.tactic_dram)

    if not self.tf32:
        config.flags &= ~(1 << self.trt.BuilderFlag.TF32.value)

    if self.sparse_weights:
        config.flags |= 1 << self.trt.BuilderFlag.SPARSE_WEIGHTS.value

    if self.strict_nans:
        config.flags |= 1 << self.trt.BuilderFlag.STRICT_NANS.value

    if self.weight_streaming:
        config.flags |= 1 << self.trt.BuilderFlag.WEIGHT_STREAMING.value

    self.configure_tactic_sources(config)
    self.configure_optimization_settings(config)

def configure_optimization_settings(self, config: trt.IBuilderConfig) -> None:
    config.builder_optimization_level = self.builder_optimization_level
    config.avg_timing_iterations = self.avg_timing_iterations

    if self.max_aux_streams is not None:
        config.max_aux_streams = self.max_aux_streams

    if self.max_num_tactics is not None:
        config.max_num_tactics = self.max_num_tactics

    if int(self.tiling_optimization_level) != 0:
        config.tiling_optimization_level = self.trt.TilingOptimizationLevel(self.tiling_optimization_level)
        config.l2_limit_for_tiling = self.l2_limit_for_tiling