feat(ml): trt-rtx (#31863)

* trt-rtx

* update dependencies

machine-learning/conftest.py

@@ -112,6 +112,7 @@ def gpus() -> Iterator[None]:
     with (
         mock.patch("immich_ml.sessions.ort._intel_gpu", return_value=Device("12.71.4-128eu", "26.22.38646.4")),
         mock.patch("immich_ml.sessions.ort._amd_gpu", return_value=Device("gfx1100", "7.2.0")),
+        mock.patch("immich_ml.sessions.ort._nvidia_gpu", return_value=Device("sm89", "617.14 10601")),
     ):
         yield
 

machine-learning/immich_ml/models/constants.py

@@ -89,6 +89,7 @@ _PADDLE_MODELS = {
 }
 
 SUPPORTED_PROVIDERS = [
+    "nv_tensorrt_rtx",
     "CUDAExecutionProvider",
     "MIGraphXExecutionProvider",
     "OpenVINOExecutionProvider",

machine-learning/immich_ml/sessions/ort.py

@@ -16,7 +16,7 @@ import numpy as np
 import onnxruntime as ort
 from immich_model.runtime import RewriteContext, RewritePlan, plan_rewrites
 from numpy.typing import NDArray
-from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf
+from onnxruntime.capi.onnxruntime_pybind11_state import Fail, InvalidProtobuf
 from pydantic import BaseModel
 
 from immich_ml.schemas import ModelInput, SessionNode, Shape
@@ -27,6 +27,17 @@ from .policy import ShapePolicy
 # the one provider that handles a free dim well; the rest miscompile it, recompile inside the run,
 # or partition around it, so they are handed one graph per shape instead
 DYNAMIC_PROVIDERS = frozenset({"CPUExecutionProvider"})
+PLANNED_AS = {"nv_tensorrt_rtx": "NvTensorRTRTXExecutionProvider"}
+
+try:
+    import onnxruntime_ep_nv_tensorrt_rtx as nv_tensorrt_rtx
+except ImportError:
+    pass
+else:
+    try:
+        ort.register_execution_provider_library(nv_tensorrt_rtx.get_ep_name(), nv_tensorrt_rtx.get_library_path())
+    except Fail as e:
+        log.info(f"TensorRT-RTX is unavailable: {e}")
 
 
 def _label(pins: Mapping[str, int]) -> str:
@@ -109,6 +120,8 @@ class GraphSpec:
                 return _intel_gpu(self.openvino_device)
             case "MIGraphXExecutionProvider":
                 return _amd_gpu(int(settings.device_id))
+            case "nv_tensorrt_rtx":
+                return _nvidia_gpu(int(settings.device_id))
         return None
 
     @cached_property
@@ -159,6 +172,8 @@ class GraphSpec:
                     options = {"arena_extend_strategy": "kSameAsRequested"}
                 case "CUDAExecutionProvider":
                     options = {"arena_extend_strategy": "kSameAsRequested", "device_id": settings.device_id}
+                case "nv_tensorrt_rtx":
+                    options = {"device_id": settings.device_id, "nv_runtime_cache_path": self.directory.as_posix()}
                 case "MIGraphXExecutionProvider":
                     options = {"device_id": settings.device_id, "migraphx_model_cache_dir": self.directory.as_posix()}
                 case "OpenVINOExecutionProvider":
@@ -249,7 +264,7 @@ def prepared(spec: GraphSpec) -> Path:
 @cache
 def _plan(provider: str) -> RewritePlan:
     version = tuple(int(piece) for piece in ort.__version__.split(".")[:3])
-    return plan_rewrites(RewriteContext(target=provider, ort_version=version))
+    return plan_rewrites(RewriteContext(target=PLANNED_AS.get(provider, provider), ort_version=version))
 
 
 @cache
@@ -355,6 +370,20 @@ def _intel_gpu(device: str) -> Device:
     raise LookupError(f"OpenCL has no GPU with the UUID of OpenVINO's {device}")
 
 
+@cache
+def _nvidia_gpu(index: int) -> Device:
+    cuda, device, major, minor = ctypes.CDLL("libcuda.so.1"), ctypes.c_int(), ctypes.c_int(), ctypes.c_int()
+    if cuda.cuInit(0) or cuda.cuDeviceGet(ctypes.byref(device), index):  # honors CUDA_VISIBLE_DEVICES
+        raise LookupError(f"CUDA has no device {index}")
+    cuda.cuDeviceGetAttribute(ctypes.byref(major), 75, device)  # CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR
+    cuda.cuDeviceGetAttribute(ctypes.byref(minor), 76, device)  # CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR
+    nvml, driver = ctypes.CDLL("libnvidia-ml.so.1"), ctypes.create_string_buffer(80)
+    nvml.nvmlInit_v2()
+    nvml.nvmlSystemGetDriverVersion(driver, 80)
+    tensorrt_rtx = ctypes.CDLL("libtensorrt_rtx.so.1").getInferLibVersion()
+    return Device(f"sm{major.value}{minor.value}", f"{driver.value.decode()} {tensorrt_rtx}")
+
+
 @cache
 def _amd_gpu(index: int) -> Device:
     """The arch MIGraphX keys its programs by, and the MIGraphX and HIP it loads, which its keys leave out."""

machine-learning/pyproject.toml

@@ -48,7 +48,10 @@ dev = ["locust>=2.15.1", { include-group = "test" }, { include-group = "lint" }]
 
 [project.optional-dependencies]
 cpu = ["onnxruntime>=1.23.2,<2"]
-cuda = ["onnxruntime-gpu>=1.27.1,<1.28"] # TODO: last ORT release supporting CUDA 12.2, also need to handle CUDA 13
+cuda = [
+  "onnxruntime-gpu>=1.27.1,<1.28", # TODO: last ORT release supporting CUDA 12.2, also need to handle CUDA 13
+  "onnxruntime-ep-nv-tensorrt-rtx-cu12>=0.4.0,<0.5 ; platform_machine == 'x86_64'",
+]
 openvino = ["onnxruntime-openvino>=1.24.1,<2"]
 armnn = ["onnxruntime>=1.23.2,<2"]
 rknn = ["onnxruntime>=1.23.2,<2", "rknn-toolkit-lite2>=2.3.0,<3"]

machine-learning/test_main.py

@@ -753,7 +753,11 @@ class TestPreparedGraphs:
     @pytest.mark.ov_device_ids(["GPU.0", "CPU"])
     @pytest.mark.parametrize(
         ("provider", "reader"),
-        [("OpenVINOExecutionProvider", "_intel_gpu"), ("MIGraphXExecutionProvider", "_amd_gpu")],
+        [
+            ("OpenVINOExecutionProvider", "_intel_gpu"),
+            ("MIGraphXExecutionProvider", "_amd_gpu"),
+            ("nv_tensorrt_rtx", "_nvidia_gpu"),
+        ],
     )
     def test_prepares_once_per_kind_of_device_and_again_for_another_version(
         self, provider: str, reader: str, ov_device_ids: mock.Mock, mocker: MockerFixture

machine-learning/uv.lock

@@ -705,6 +705,7 @@ cpu = [
     { name = "onnxruntime" },
 ]
 cuda = [
+    { name = "onnxruntime-ep-nv-tensorrt-rtx-cu12", marker = "platform_machine == 'x86_64'" },
     { name = "onnxruntime-gpu" },
 ]
 openvino = [
@@ -769,6 +770,7 @@ requires-dist = [
     { name = "onnxruntime", marker = "extra == 'armnn'", specifier = ">=1.23.2,<2" },
     { name = "onnxruntime", marker = "extra == 'cpu'", specifier = ">=1.23.2,<2" },
     { name = "onnxruntime", marker = "extra == 'rknn'", specifier = ">=1.23.2,<2" },
+    { name = "onnxruntime-ep-nv-tensorrt-rtx-cu12", marker = "platform_machine == 'x86_64' and extra == 'cuda'", specifier = ">=0.4.0,<0.5" },
     { name = "onnxruntime-gpu", marker = "extra == 'cuda'", specifier = ">=1.27.1,<1.28", index = "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/" },
     { name = "onnxruntime-migraphx", marker = "extra == 'rocm'", specifier = ">=1.23.2,<2" },
     { name = "onnxruntime-openvino", marker = "extra == 'openvino'", specifier = ">=1.24.1,<2" },
@@ -1346,6 +1348,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/91/f0/40f74b7c00077e1e25627067ed98a70df1fef5c0e21b82849190d312554e/onnxruntime-1.30.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:67ad7f03433b6462c627d0f555dece80e6a26bc71e8542ced35cebd32142d1b7", size = 23579340, upload-time = "2026-09-10T16:32:15.532Z" },
 ]
 
+[[package]]
+name = "onnxruntime-ep-nv-tensorrt-rtx-cu12"
+version = "0.4.0"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/52/a6/1405069d40e2d6d37b21031e5dc2cc8b8197fbb2e3d4e4f65b9155c57e4f/onnxruntime_ep_nv_tensorrt_rtx_cu12-0.4.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:221807b1797a6270f37dfd7981f66e25ca7483f3883bfe23ef2b6a4198fe2b10", size = 167721016, upload-time = "2026-08-14T08:07:42.23Z" },
+    { url = "https://files.pythonhosted.org/packages/8b/e4/f19fdbffcf8faf18d9f81973f81e3d378ef371f857b049a3f355535a015c/onnxruntime_ep_nv_tensorrt_rtx_cu12-0.4.0-py3-none-win_amd64.whl", hash = "sha256:5d9dba6aafd8e0863f34d0e5f51ad7c7ae9a5d6a92a55fe45d9637d3a5f1e1a7", size = 105822356, upload-time = "2026-08-12T11:13:34.401Z" },
+]
+
 [[package]]
 name = "onnxruntime-gpu"
 version = "1.27.1"