From f553bf0067264b96273b138df57ed39865c1ba38 Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Sat, 7 Jun 2025 11:07:18 -0600
Subject: [PATCH] Add missing packages and cleanup ldconfig

---
 docker/tensorrt/Dockerfile.amd64                | 17 +++++++++++++----
 .../rootfs/etc/ld.so.conf.d/cuda_tensorrt.conf  |  7 +++----
 docker/tensorrt/requirements-amd64.txt          |  4 +++-
 docker/tensorrt/trt.hcl                         |  3 ++-
 4 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/docker/tensorrt/Dockerfile.amd64 b/docker/tensorrt/Dockerfile.amd64
index 3ae927018..8e0d756f6 100644
--- a/docker/tensorrt/Dockerfile.amd64
+++ b/docker/tensorrt/Dockerfile.amd64
@@ -11,12 +11,21 @@ ARG PIP_BREAK_SYSTEM_PACKAGES
 
 # Install TensorRT wheels
 COPY docker/tensorrt/requirements-amd64.txt /requirements-tensorrt.txt
-RUN pip3 uninstall -y onnxruntime-openvino \
-    && pip3 install -U -r /requirements-tensorrt.txt \
-    && ldconfig
+COPY docker/main/requirements-wheels.txt /requirements-wheels.txt
+RUN  pip3 wheel --wheel-dir=/trt-wheels -c /requirements-wheels.txt -r /requirements-tensorrt.txt
+
+FROM deps AS frigate-tensorrt
+ARG PIP_BREAK_SYSTEM_PACKAGES
+
+RUN --mount=type=bind,from=trt-wheels,source=/trt-wheels,target=/deps/trt-wheels \
+    pip3 uninstall -y onnxruntime-openvino \
+    && pip3 install -U /deps/trt-wheels/*.whl
+
+COPY --from=rootfs / /
+COPY docker/tensorrt/detector/rootfs/etc/ld.so.conf.d /etc/ld.so.conf.d
+RUN ldconfig
 
 WORKDIR /opt/frigate/
-COPY --from=rootfs / /
 
 # Dev Container w/ TRT
 FROM devcontainer AS devcontainer-trt
diff --git a/docker/tensorrt/detector/rootfs/etc/ld.so.conf.d/cuda_tensorrt.conf b/docker/tensorrt/detector/rootfs/etc/ld.so.conf.d/cuda_tensorrt.conf
index f66af7dc6..b00d4b4e4 100644
--- a/docker/tensorrt/detector/rootfs/etc/ld.so.conf.d/cuda_tensorrt.conf
+++ b/docker/tensorrt/detector/rootfs/etc/ld.so.conf.d/cuda_tensorrt.conf
@@ -1,7 +1,6 @@
-/usr/local/lib
-/usr/local/cuda
-/usr/local/lib/python3.11/dist-packages/tensorrt
 /usr/local/lib/python3.11/dist-packages/nvidia/cudnn/lib
 /usr/local/lib/python3.11/dist-packages/nvidia/cuda_runtime/lib
 /usr/local/lib/python3.11/dist-packages/nvidia/cublas/lib
-/usr/local/lib/python3.11/dist-packages/nvidia/cufft/lib
\ No newline at end of file
+/usr/local/lib/python3.11/dist-packages/nvidia/cufft/lib
+/usr/local/lib/python3.11/dist-packages/nvidia/curand/lib/
+/usr/local/lib/python3.11/dist-packages/nvidia/cuda_nvrtc/lib/
\ No newline at end of file
diff --git a/docker/tensorrt/requirements-amd64.txt b/docker/tensorrt/requirements-amd64.txt
index a53e3cc6e..615478549 100644
--- a/docker/tensorrt/requirements-amd64.txt
+++ b/docker/tensorrt/requirements-amd64.txt
@@ -1,10 +1,12 @@
 # NVidia TensorRT Support (amd64 only)
 --extra-index-url 'https://pypi.nvidia.com'
-numpy < 1.24; platform_machine == 'x86_64'
 cython == 3.0.*; platform_machine == 'x86_64'
 nvidia-cuda-runtime-cu12 == 12.1.*; platform_machine == 'x86_64'
+nvidia-cublas-cu12==12.6.4.*; platform_machine == 'x86_64'
 nvidia-cudnn-cu12 == 9.5.0.*; platform_machine == 'x86_64'
 nvidia-cufft-cu12==11.*; platform_machine == 'x86_64'
+nvidia-curand-cu12==10.3.9.*; platform_machine == 'x86_64'
+nvidia-cuda-nvrtc-cu12==12.6.*; platform_machine == 'x86_64'
 onnx==1.16.*; platform_machine == 'x86_64'
 onnxruntime-gpu==1.22.*; platform_machine == 'x86_64'
 protobuf==3.20.3; platform_machine == 'x86_64'
diff --git a/docker/tensorrt/trt.hcl b/docker/tensorrt/trt.hcl
index 80757ba6d..8f631b6aa 100644
--- a/docker/tensorrt/trt.hcl
+++ b/docker/tensorrt/trt.hcl
@@ -93,7 +93,8 @@ target "tensorrt" {
   context = "."
   contexts = {
     wget = "target:wget",
-    tensorrt-base = "target:tensorrt-base",
+    wheels = "target:wheels",
+    deps = "target:deps",
     rootfs = "target:rootfs"
   }
   target = "frigate-tensorrt"