From d244a2532cf9c93481e3edf2f742bce5c9f28f48 Mon Sep 17 00:00:00 2001
From: Nate Meyer <nate.devel@gmail.com>
Date: Tue, 20 Jun 2023 00:21:42 -0400
Subject: [PATCH] Update tensorrt_models script to convert models from the
 frigate container

---
 docker/install_deps.sh                        |  2 +-
 .../etc/ld.so.conf.d/cuda_tensorrt.conf       |  1 +
 docker/tensorrt_models.sh                     | 24 +++++++++----------
 requirements-tensorrt.txt                     |  8 +++++--
 4 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/docker/install_deps.sh b/docker/install_deps.sh
index 25b6951b5..7d5242d83 100755
--- a/docker/install_deps.sh
+++ b/docker/install_deps.sh
@@ -68,7 +68,7 @@ if [[ "${TARGETARCH}" == "arm64" ]]; then
         libva-drm2 mesa-va-drivers
 fi
 
-apt-get purge gnupg apt-transport-https wget xz-utils -y
+apt-get purge gnupg apt-transport-https xz-utils -y
 apt-get clean autoclean -y
 apt-get autoremove --purge -y
 rm -rf /var/lib/apt/lists/*
diff --git a/docker/rootfs/etc/ld.so.conf.d/cuda_tensorrt.conf b/docker/rootfs/etc/ld.so.conf.d/cuda_tensorrt.conf
index d4248d047..fe16ed9c5 100644
--- a/docker/rootfs/etc/ld.so.conf.d/cuda_tensorrt.conf
+++ b/docker/rootfs/etc/ld.so.conf.d/cuda_tensorrt.conf
@@ -1,3 +1,4 @@
+/usr/local/lib
 /usr/local/lib/python3.9/dist-packages/nvidia/cudnn/lib
 /usr/local/lib/python3.9/dist-packages/nvidia/cuda_runtime/lib
 /usr/local/lib/python3.9/dist-packages/nvidia/cublas/lib
diff --git a/docker/tensorrt_models.sh b/docker/tensorrt_models.sh
index 5c09a79fe..3455e8038 100755
--- a/docker/tensorrt_models.sh
+++ b/docker/tensorrt_models.sh
@@ -2,33 +2,31 @@
 
 set -euxo pipefail
 
-CUDA_HOME=/usr/local/cuda
-LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64
-OUTPUT_FOLDER=/tensorrt_models
+OUTPUT_FOLDER=/trt_models
 echo "Generating the following TRT Models: ${YOLO_MODELS:="yolov4-tiny-288,yolov4-tiny-416,yolov7-tiny-416"}"
 
 # Create output folder
 mkdir -p ${OUTPUT_FOLDER}
 
-# Install packages
-pip install --upgrade pip && pip install onnx==1.15.0 protobuf==3.20.3
-
 # Clone tensorrt_demos repo
-git clone --depth 1 https://github.com/yeahme49/tensorrt_demos.git /tensorrt_demos
+# git clone --depth 1 https://github.com/NateMeyer/tensorrt_demos.git -b conditional_download /tmp/tensorrt_demos
+cd /tmp/ && wget -qO tensorrt_demos.zip https://github.com/NateMeyer/tensorrt_demos/archive/refs/heads/conditional_download.zip
+unzip tensorrt_demos.zip
 
-# Build libyolo
-cd /tensorrt_demos/plugins && make all
-cp libyolo_layer.so ${OUTPUT_FOLDER}/libyolo_layer.so
+cp /usrl/local/lib/libyolo_layer.so /tmp/tensorrt_demos-conditional_download/plugins/libyolo_layer.so
 
 # Download yolo weights
-cd /tensorrt_demos/yolo && ./download_yolo.sh
+cd /tmp/tensorrt_demos-conditional_download/yolo && ./download_yolo.sh $YOLO_MODELS
 
 # Build trt engine
-cd /tensorrt_demos/yolo
+cd /tmp/tensorrt_demos-conditional_download/yolo
 
 for model in ${YOLO_MODELS//,/ }
 do
     python3 yolo_to_onnx.py -m ${model}
     python3 onnx_to_tensorrt.py -m ${model}
-    cp /tensorrt_demos/yolo/${model}.trt ${OUTPUT_FOLDER}/${model}.trt;
+    cp /tmp/tensorrt_demos-conditional_download/yolo/${model}.trt ${OUTPUT_FOLDER}/${model}.trt;
 done
+
+# Cleanup repo
+rm -r /tmp/tensorrt_demos-conditional_download
diff --git a/requirements-tensorrt.txt b/requirements-tensorrt.txt
index 2ce288fdc..3b6d099a3 100644
--- a/requirements-tensorrt.txt
+++ b/requirements-tensorrt.txt
@@ -1,9 +1,13 @@
 # NVidia TensorRT Support (amd64 only)
-nvidia-pyindex; platform_machine == 'x86_64'
+--extra-index-url 'https://pypi.nvidia.com'
 tensorrt == 8.6.1; platform_machine == 'x86_64'
+tensorrt-libs == 8.6.1; platform_machine == 'x86_64'
+tensorrt-bindings == 8.6.1; platform_machine == 'x86_64'
 cuda-python == 12.1; platform_machine == 'x86_64'
 cython == 0.29.*; platform_machine == 'x86_64'
 nvidia-cuda-runtime-cu12 == 12.1.*; platform_machine == 'x86_64'
 nvidia-cublas-cu12 == 12.1.*; platform_machine == 'x86_64'
 nvidia-cudnn-cu12 == 8.9.*; platform_machine == 'x86_64'
-nvidia-cuda-nvrtc-cu12 == 12.1.*; platform_machine == 'x86_64'
\ No newline at end of file
+nvidia-cuda-nvrtc-cu12 == 12.1.*; platform_machine == 'x86_64'
+onnx==1.14.0; platform_machine == 'x86_64'
+protobuf==3.20.3; platform_machine == 'x86_64'
\ No newline at end of file