add docker setup for tensorrt-llm; update readme

README.md

@@ -100,12 +100,17 @@ This command streams audio into the server from a HLS stream. It uses the same o
  docker run -it --gpus all -p 9090:9090 whisper-live:latest
 ```
 
+- TensorRT.Follow [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md) in order to setup and use TensorRT backend. 
+```bash
+ docker build --file docker/Dockerfile.tensorrt --tag tensorrt_llm/devel:latest .
+```
+
 - CPU
 ```bash
  docker build . -t whisper-live -f docker/Dockerfile.cpu
  docker run -it -p 9090:9090 whisper-live:latest
 ```
-**Note**: This only builds the docker image for `faster_whisper` backend. Follow [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md) in order to setup and use TensorRT backend. By default we use "small" model size. To build docker image for a different model size, change the size in server.py and then build the docker image.
+**Note**: By default we use "small" model size. To build docker image for a different model size, change the size in server.py and then build the docker image.
 
 ## Future Work
 - [ ] Add translation to other languages on top of transcription.

TensorRT_whisper.md

@@ -4,27 +4,21 @@ We have only tested the TensorRT backend in docker so, we recommend docker for a
 ## Installation
 - Install [docker](https://docs.docker.com/engine/install/)
 - Install [nvidia-container-toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
-- Pull the pytorch docker image.
-```bash
-docker pull nvcr.io/nvidia/pytorch_23.10-py3
-```
 - Clone this repo.
 ```bash
 git clone https://github.com/collabora/WhisperLive.git
+cd WhisperLive
 ```
+- Build the TensorRT-LLM docker image
+```bash
+docker build --file docker/Dockerfile.tensorrt --tag tensorrt_llm/devel:latest .
+```
+**NOTE**: This could take some time.
 - Next, we run the docker image and mount WhisperLive repo to the containers `/home` directory.
 ```bash
-docker run -it --gpus all --shm-size=64g /path/to/WhisperLive:/home/WhisperLive nvcr.io/nvidia/pytorch_23.10-py3
+docker run -it --gpus all --shm-size=64g /path/to/WhisperLive:/home/WhisperLive tensorrt_llm/devel:latest
 ```
-- Build `tensorrt-llm`.
+- Once inside the docker container, make sure to test the installation.
-```bash
-cd /home/
-cp WhisperLive/scripts/install_tensorrt_llm.sh .
-bash install_tensorrt_llm.sh
-```
-This should clone the [NVIDIA/TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) and build it as well.
-
-- Test the installation.
 ```bash
 export ENV=${ENV:-/etc/shinit_v2}
 source $ENV
@@ -36,9 +30,6 @@ python -c "import torch; import tensorrt; import tensorrt_llm"
 ```bash
 cd /home/TensorRT-LLM/examples/whisper
 ``` 
-- Currently, by default TensorRT-LLM only supports `large-v2` and `large-v3`. In this repo, we use `small.en`.
-
-- Edit `build.py` to support all the model sizes i.e. `["tiny", "tiny.en", "base", "base.en", "small", "small.en", "medium", "medium.en"]`. In order to do that, update the list [`choices`](https://github.com/NVIDIA/TensorRT-LLM/blob/a75618df24e97ecf92b8899ca3c229c4b8097dda/examples/whisper/build.py#L58) with the model size you prefer for your WhisperLive server.
 
 - Download the models from [here](https://github.com/openai/whisper/blob/ba3f3cd54b0e5b8ce1ab3de13e32122d0d5f98ab/whisper/__init__.py#L17C1-L30C2)
 ```bash
@@ -83,4 +74,3 @@ python3 run_server.py --port 9090 \
                       --whisper_tensorrt_path /home/TensorRT-LLM/examples/whisper/whisper_small_en \
                       --trt_multilingual
 ```
-

docker/Dockerfile.tensorrt

@@ -0,0 +1,8 @@
+ARG BASE_IMAGE=nvcr.io/nvidia/pytorch
+ARG BASE_TAG=23.10-py3
+
+FROM ${BASE_IMAGE}:${BASE_TAG} as base
+
+WORKDIR /home
+COPY script/install_trt_llm.sh install_trt_llm.sh
+RUN bash install_trt_llm.sh && rm install_trt_llm.sh

scripts/install_tensorrt_llm.sh

@@ -29,9 +29,9 @@ echo 'export PATH=$PATH:/usr/local/bin' >> "${ENV}"
 echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib' >> "${ENV}"
 source $ENV
 
-cd /root/TensorRT-LLM/docker/common/
+cd /home/TensorRT-LLM/docker/common/
 bash install_mpi4py.sh
 source $ENV
 
-cd /root/TensorRT-LLM
+cd /home/TensorRT-LLM
 python3 scripts/build_wheel.py --clean --trt_root /usr/local/tensorrt