diff --git a/README.md b/README.md index 9db1d37..16be1a8 100644 --- a/README.md +++ b/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. diff --git a/TensorRT_whisper.md b/TensorRT_whisper.md index f9a0042..6fcda77 100644 --- a/TensorRT_whisper.md +++ b/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`. -```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. +- Once inside the docker container, make sure to 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 ``` - diff --git a/docker/Dockerfile.tensorrt b/docker/Dockerfile.tensorrt new file mode 100644 index 0000000..9bcf7b0 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/scripts/install_tensorrt_llm.sh b/scripts/install_tensorrt_llm.sh index 4463dfa..b0340a6 100644 --- a/scripts/install_tensorrt_llm.sh +++ b/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