add eos to faster_whisper server

Signed-off-by: makaveli10 <vineet.suryan@collabora.com>

.github/workflows/ci.yml

@@ -99,6 +99,35 @@ jobs:
           push: true
           tags: ghcr.io/collabora/whisperlive-cpu:latest
   
+  build-and-push-docker-tensorrt:
+    needs: [run-tests, check-code-format]
+    timeout-minutes: 20
+    runs-on: ubuntu-22.04
+    if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/'))
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v1
+        with:
+          registry: ghcr.io
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.GHCR_TOKEN }}
+
+      - name: Docker Prune
+        run: docker system prune -af
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+
+      - name: Build and push Docker GPU image
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          file: docker/Dockerfile.tensorrt
+          push: true
+          tags: ghcr.io/collabora/whisperlive-tensorrt:latest
+
   build-and-push-docker-gpu:
     needs: [run-tests, check-code-format, build-and-push-docker-cpu]
     timeout-minutes: 20

README.md

@@ -36,7 +36,7 @@ python3 run_server.py --port 9090 \
   
 # running with custom model
 python3 run_server.py --port 9090 \
-                      --backend faster_whisper
+                      --backend faster_whisper \
                       -fw "/path/to/custom/faster/whisper/model"
 ```
 
@@ -53,10 +53,30 @@ python3 run_server.py -p 9090 \
                       -trt /home/TensorRT-LLM/examples/whisper/whisper_small \
                       -m
 ```
+#### Controlling OpenMP Threads
+To control the number of threads used by OpenMP, you can set the `OMP_NUM_THREADS` environment variable. This is useful for managing CPU resources and ensuring consistent performance. If not specified, `OMP_NUM_THREADS` is set to `1` by default. You can change this by using the `--omp_num_threads` argument:
+```bash
+python3 run_server.py --port 9090 \
+                      --backend faster_whisper \
+                      --omp_num_threads 4
+```
+
+#### Single model mode
+By default, when running the server without specifying a model, the server will instantiate a new whisper model for every client connection. This has the advantage, that the server can use different model sizes, based on the client's requested model size. On the other hand, it also means you have to wait for the model to be loaded upon client connection and you will have increased (V)RAM usage.
+
+When serving a custom TensorRT model using the `-trt` or a custom faster_whisper model using the `-fw` option, the server will instead only instantiate the custom model once and then reuse it for all client connections.
+
+If you don't want this, set `--no_single_model`.
 
 
 ### Running the Client
-- Initializing the client:
+- Initializing the client with below parameters:
+  - `lang`: Language of the input audio, applicable only if using a multilingual model.
+  - `translate`: If set to `True` then translate from any language to `en`.
+  - `model`: Whisper model size.
+  - `use_vad`: Whether to use `Voice Activity Detection` on the server.
+  - `save_output_recording`: Set to True to save the microphone input as a `.wav` file during live transcription. This option is helpful for recording sessions for later playback or analysis. Defaults to `False`. 
+  - `output_recording_filename`: Specifies the `.wav` file path where the microphone input will be saved if `save_output_recording` is set to `True`.
 ```python
 from whisper_live.client import TranscriptionClient
 client = TranscriptionClient(
@@ -66,11 +86,13 @@ client = TranscriptionClient(
   translate=False,
   model="small",
   use_vad=False,
+  save_output_recording=True,                         # Only used for microphone input, False by Default
+  output_recording_filename="./output_recording.wav"  # Only used for microphone input
 )
 ```
 It connects to the server running on localhost at port 9090. Using a multilingual model, language for the transcription will be automatically detected. You can also use the language option to specify the target language for the transcription, in this case, English ("en"). The translate option should be set to `True` if we want to translate from the source language to English and `False` if we want to transcribe in the source language.
 
-- Trancribe an audio file:
+- Transcribe an audio file:
 ```python
 client("tests/jfk.wav")
 ```
@@ -80,9 +102,14 @@ client("tests/jfk.wav")
 client()
 ```
 
+- To transcribe from a RTSP stream:
+```python
+client(rtsp_url="rtsp://admin:admin@192.168.0.1/rtsp")
+```
+
 - To transcribe from a HLS stream:
 ```python
-client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/bbc_1xtra.isml/bbc_1xtra-audio%3d96000.norewind.m3u8") 
+client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/bbc_1xtra.isml/bbc_1xtra-audio%3d96000.norewind.m3u8")
 ```
 
 ## Browser Extensions
@@ -96,7 +123,18 @@ client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/b
   docker run -it --gpus all -p 9090:9090 ghcr.io/collabora/whisperlive-gpu:latest
   ```
 
-  - TensorRT. Follow [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md) in order to setup docker and use TensorRT backend. We provide a pre-built docker image which has TensorRT-LLM built and ready to use.
+  - TensorRT. 
+  ```bash
+  docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt
+
+  # Build tiny.en engine
+  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en
+
+  # Run server with tiny.en
+  python3 run_server.py --port 9090 \
+                        --backend tensorrt \
+                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en"
+  ```
 
 - CPU
 ```bash

TensorRT_whisper.md

@@ -1,6 +1,6 @@
-# Whisper-TensorRT
+# WhisperLive-TensorRT
 We have only tested the TensorRT backend in docker so, we recommend docker for a smooth TensorRT backend setup.
-**Note**: We use [our fork to setup TensorRT](https://github.com/makaveli10/TensorRT-LLM)
+**Note**: We use `tensorrt_llm==0.9.0`
 
 ## Installation
 - Install [docker](https://docs.docker.com/engine/install/)
@@ -12,56 +12,31 @@ git clone https://github.com/collabora/WhisperLive.git
 cd WhisperLive
 ```
 
-- Pull the TensorRT-LLM docker image which we prebuilt for WhisperLive TensorRT backend.
+- Run WhisperLive TensorRT in docker
 ```bash
-docker pull ghcr.io/collabora/whisperbot-base:latest
+docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt:latest
 ```
 
-- Next, we run the docker image and mount WhisperLive repo to the containers `/home` directory.
-```bash
-docker run -it --gpus all --shm-size=8g \
-       --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \
-       -p 9090:9090 -v /path/to/WhisperLive:/home/WhisperLive \
-       ghcr.io/collabora/whisperbot-base:latest
-```
-
-- Make sure to test the installation. 
-```bash
-# export ENV=${ENV:-/etc/shinit_v2} 
-# source $ENV
-python -c "import torch; import tensorrt; import tensorrt_llm"
-```
-**NOTE**: Uncomment and update library paths if imports fail.
-
 ## Whisper TensorRT Engine
-- We build `small.en` and `small` multilingual TensorRT engine. The script logs the path of the directory with Whisper TensorRT engine. We need the model_path to run the server.
+- We build `small.en` and `small` multilingual TensorRT engine as examples below. The script logs the path of the directory with Whisper TensorRT engine. We need that model_path to run the server.
 ```bash
 # convert small.en
-bash scripts/build_whisper_tensorrt.sh /root/TensorRT-LLM-examples small.en
+bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en
 
 # convert small multilingual model
-bash scripts/build_whisper_tensorrt.sh /root/TensorRT-LLM-examples small
+bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small
 ```
 
 ## Run WhisperLive Server with TensorRT Backend
 ```bash
-cd /home/WhisperLive
-
-# Install requirements
-apt update && bash scripts/setup.sh
-pip install -r requirements/server.txt
-
-# Required to create mel spectogram
-wget --directory-prefix=assets assets/mel_filters.npz https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz
-
 # Run English only model
 python3 run_server.py --port 9090 \
                       --backend tensorrt \
-                      --trt_model_path "path/to/whisper_trt/from/build/step"
+                      --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en"
 
 # Run Multilingual model
 python3 run_server.py --port 9090 \
                       --backend tensorrt \
-                      --trt_model_path "path/to/whisper_trt/from/build/step" \
+                      --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small" \
                       --trt_multilingual
 ```

docker/Dockerfile.cpu

@@ -1,22 +1,23 @@
-FROM python:3.8-slim-buster
+FROM python:3.10-bookworm
 
 ARG DEBIAN_FRONTEND=noninteractive
 
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    curl \
-    ca-certificates \
-    sudo \
-    git \
-    bzip2 \
-    libx11-6 \
- && apt-get clean \
- && rm -rf /var/lib/apt/lists/*
+# install lib required for pyaudio
+RUN apt update && apt install -y portaudio19-dev && apt-get clean && rm -rf /var/lib/apt/lists/*
 
+# update pip to support for whl.metadata -> less downloading
+RUN pip install --no-cache-dir -U "pip>=24"
+
+# create a working directory
+RUN mkdir /app
 WORKDIR /app
 
-COPY scripts/setup.sh requirements/server.txt /app/
+# install pytorch, but without the nvidia-libs that are only necessary for gpu
+RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu
 
-RUN apt update && bash setup.sh && pip install -r server.txt
+# install the requirements for running the whisper-live server
+COPY requirements/server.txt /app/
+RUN pip install --no-cache-dir -r server.txt && rm server.txt
 
 COPY whisper_live /app/whisper_live
 COPY run_server.py /app

docker/Dockerfile.gpu

@@ -1,33 +1,26 @@
-FROM nvidia/cuda:12.2.2-cudnn8-runtime-ubuntu22.04
+FROM python:3.10-bookworm
+
 ARG DEBIAN_FRONTEND=noninteractive
 
-# Remove any third-party apt sources to avoid issues with expiring keys.
-RUN rm -f /etc/apt/sources.list.d/*.list
+# install lib required for pyaudio
+RUN apt update && apt install -y portaudio19-dev && apt-get clean && rm -rf /var/lib/apt/lists/*
 
-# Install some basic utilities.
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    curl \
-    ca-certificates \
-    sudo \
-    git \
-    bzip2 \
-    libx11-6 \
-    python3-dev \
-    python3-pip \
-    && python3 -m pip install --upgrade pip \
-    && rm -rf /var/lib/apt/lists/*
+# update pip to support for whl.metadata -> less downloading
+RUN pip install --no-cache-dir -U "pip>=24"
 
-# Create a working directory.
+# create a working directory
 RUN mkdir /app
 WORKDIR /app
 
-COPY scripts/setup.sh requirements/server.txt /app
+# install the requirements for running the whisper-live server
+COPY requirements/server.txt /app/
+RUN pip install --no-cache-dir -r server.txt && rm server.txt
 
-RUN apt update && bash setup.sh && rm setup.sh
-RUN pip install -r server.txt && rm server.txt
+# make the paths of the nvidia libs installed as wheels visible. equivalent to:
+# export LD_LIBRARY_PATH=`python3 -c 'import os; import nvidia.cublas.lib; import nvidia.cudnn.lib; print(os.path.dirname(nvidia.cublas.lib.__file__) + ":" + os.path.dirname(nvidia.cudnn.lib.__file__))'`
+ENV LD_LIBRARY_PATH="/usr/local/lib/python3.10/site-packages/nvidia/cublas/lib:/usr/local/lib/python3.10/site-packages/nvidia/cudnn/lib"
 
 COPY whisper_live /app/whisper_live
-
 COPY run_server.py /app
 
-CMD ["python3", "run_server.py"]
+CMD ["python", "run_server.py"]

docker/Dockerfile.tensorrt

@@ -0,0 +1,28 @@
+FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update && apt-get install -y \
+    python3.10 python3-pip openmpi-bin libopenmpi-dev git wget \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN pip3 install --no-cache-dir -U tensorrt_llm==0.9.0 --extra-index-url https://pypi.nvidia.com
+
+WORKDIR /app
+
+RUN git clone -b v0.9.0 --depth 1 https://github.com/NVIDIA/TensorRT-LLM.git && \
+    mv TensorRT-LLM/examples ./TensorRT-LLM-examples && \
+    rm -rf TensorRT-LLM
+
+COPY assets/ ./assets
+RUN wget -nc -P assets/ https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz
+
+COPY scripts/setup.sh ./
+RUN apt update && bash setup.sh && rm setup.sh
+
+COPY requirements/server.txt .
+RUN pip install --no-cache-dir -r server.txt && rm server.txt
+
+COPY whisper_live ./whisper_live
+COPY scripts/build_whisper_tensorrt.sh .
+COPY run_server.py .

requirements/server.txt

@@ -9,4 +9,5 @@ soundfile
 ffmpeg-python
 scipy
 jiwer
-evaluate
+evaluate
+numpy<2

run_server.py

@@ -1,5 +1,5 @@
 import argparse
-from whisper_live.server import TranscriptionServer
+import os
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -21,12 +21,23 @@ if __name__ == "__main__":
     parser.add_argument('--trt_multilingual', '-m',
                         action="store_true",
                         help='Boolean only for TensorRT model. True if multilingual.')
+    parser.add_argument('--omp_num_threads', '-omp',
+                        type=int,
+                        default=1,
+                        help="Number of threads to use for OpenMP")
+    parser.add_argument('--no_single_model', '-nsm',
+                        action='store_true',
+                        help='Set this if every connection should instantiate its own model. Only relevant for custom model, passed using -trt or -fw.')
     args = parser.parse_args()
 
     if args.backend == "tensorrt":
         if args.trt_model_path is None:
             raise ValueError("Please Provide a valid tensorrt model path")
 
+    if "OMP_NUM_THREADS" not in os.environ:
+        os.environ["OMP_NUM_THREADS"] = str(args.omp_num_threads)
+
+    from whisper_live.server import TranscriptionServer
     server = TranscriptionServer()
     server.run(
         "0.0.0.0",
@@ -34,5 +45,6 @@ if __name__ == "__main__":
         backend=args.backend,
         faster_whisper_custom_model_path=args.faster_whisper_custom_model_path,
         whisper_tensorrt_path=args.trt_model_path,
-        trt_multilingual=args.trt_multilingual
+        trt_multilingual=args.trt_multilingual,
+        single_model=not args.no_single_model,
     )

scripts/build_whisper_tensorrt.sh

@@ -57,7 +57,7 @@ download_and_build_model() {
     local output_dir="whisper_${model_name//./_}"
     echo "$output_dir"
     echo "Running build script for $model_name with output directory $output_dir"
-    python3 build.py --output_dir "$output_dir" --use_gpt_attention_plugin --use_gemm_plugin --use_bert_attention_plugin --model_name "$model_name"
+    python3 build.py --output_dir "$output_dir" --use_gpt_attention_plugin --use_gemm_plugin  --use_bert_attention_plugin --enable_context_fmha --model_name "$model_name"
     echo "Whisper $model_name TensorRT engine built."
     echo "========================================="
     echo "Model is located at: $(pwd)/$output_dir"

whisper_live/__version__.py

@@ -1 +1 @@
-__version__ = "0.4.1"
+__version__ = "0.5.0"

whisper_live/client.py

@@ -1,4 +1,5 @@
 import os
+import shutil
 import wave
 
 import numpy as np
@@ -259,6 +260,7 @@ class Client:
         while time.time() - self.last_response_received < self.disconnect_if_no_response_for:
             continue
 
+
 class TranscriptionTeeClient:
     """
     Client for handling audio recording, streaming, and transcription tasks via one or more
@@ -272,7 +274,7 @@ class TranscriptionTeeClient:
     Attributes:
         clients (list): the underlying Client instances responsible for handling WebSocket connections.
     """
-    def __init__(self, clients):
+    def __init__(self, clients, save_output_recording=False, output_recording_filename="./output_recording.wav"):
         self.clients = clients
         if not self.clients:
             raise Exception("At least one client is required.")
@@ -281,6 +283,8 @@ class TranscriptionTeeClient:
         self.channels = 1
         self.rate = 16000
         self.record_seconds = 60000
+        self.save_output_recording = save_output_recording
+        self.output_recording_filename = output_recording_filename
         self.frames = b""
         self.p = pyaudio.PyAudio()
         try:
@@ -295,7 +299,7 @@ class TranscriptionTeeClient:
             print(f"[WARN]: Unable to access microphone. {error}")
             self.stream = None
 
-    def __call__(self, audio=None, hls_url=None):
+    def __call__(self, audio=None, rtsp_url=None, hls_url=None, save_file=None):
         """
         Start the transcription process.
 
@@ -307,6 +311,10 @@ class TranscriptionTeeClient:
             audio (str, optional): Path to an audio file for transcription. Default is None, which triggers live recording.
 
         """
+        assert sum(
+            source is not None for source in [audio, rtsp_url, hls_url]
+        ) <= 1, 'You must provide only one selected source'
+
         print("[INFO]: Waiting for server ready ...")
         for client in self.clients:
             while not client.recording:
@@ -316,10 +324,12 @@ class TranscriptionTeeClient:
 
         print("[INFO]: Server Ready!")
         if hls_url is not None:
-            self.process_hls_stream(hls_url)
+            self.process_hls_stream(hls_url, save_file)
         elif audio is not None:
             resampled_file = utils.resample(audio)
             self.play_file(resampled_file)
+        elif rtsp_url is not None:
+            self.process_rtsp_stream(rtsp_url)
         else:
             self.record()
 
@@ -398,25 +408,30 @@ class TranscriptionTeeClient:
                 self.write_all_clients_srt()
                 print("[INFO]: Keyboard interrupt.")
 
-    def process_hls_stream(self, hls_url):
+    def process_rtsp_stream(self, rtsp_url):
+        """
+        Connect to an RTSP source, process the audio stream, and send it for trascription.
+
+        Args:
+            rtsp_url (str): The URL of the RTSP stream source.
+        """
+        process = self.get_rtsp_ffmpeg_process(rtsp_url)
+        self.handle_ffmpeg_process(process, stream_type='RTSP')
+
+    def process_hls_stream(self, hls_url, save_file):
         """
         Connect to an HLS source, process the audio stream, and send it for transcription.
 
         Args:
             hls_url (str): The URL of the HLS stream source.
+            save_file （str, optional): Local path to save the network stream.
         """
-        print("[INFO]: Connecting to HLS stream...")
-        process = None  # Initialize process to None
+        process = self.get_hls_ffmpeg_process(hls_url, save_file)
+        self.handle_ffmpeg_process(process, stream_type='HLS')
 
+    def handle_ffmpeg_process(self, process, stream_type):
+        print(f"[INFO]: Connecting to {stream_type} stream...")
         try:
-            # Connecting to the HLS stream using ffmpeg-python
-            process = (
-                ffmpeg
-                .input(hls_url, threads=0)
-                .output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
-                .run_async(pipe_stdout=True, pipe_stderr=True)
-            )
-
             # Process the stream
             while True:
                 in_bytes = process.stdout.read(self.chunk * 2)  # 2 bytes per sample
@@ -426,14 +441,79 @@ class TranscriptionTeeClient:
                 self.multicast_packet(audio_array.tobytes())
 
         except Exception as e:
-            print(f"[ERROR]: Failed to connect to HLS stream: {e}")
+            print(f"[ERROR]: Failed to connect to {stream_type} stream: {e}")
         finally:
+            self.close_all_clients()
+            self.write_all_clients_srt()
             if process:
                 process.kill()
 
-        print("[INFO]: HLS stream processing finished.")
+        print(f"[INFO]: {stream_type} stream processing finished.")
 
-    def record(self, out_file="output_recording.wav"):
+    def get_rtsp_ffmpeg_process(self, rtsp_url):
+        return (
+            ffmpeg
+            .input(rtsp_url, threads=0)
+            .output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
+            .run_async(pipe_stdout=True, pipe_stderr=True)
+        )
+
+    def get_hls_ffmpeg_process(self, hls_url, save_file):
+        if save_file is None:
+            process = (
+                ffmpeg
+                .input(hls_url, threads=0)
+                .output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
+                .run_async(pipe_stdout=True, pipe_stderr=True)
+            )
+        else:
+            input = ffmpeg.input(hls_url, threads=0)
+            output_file = input.output(save_file, acodec='copy', vcodec='copy').global_args('-loglevel', 'quiet')
+            output_std = input.output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
+            process = (
+                ffmpeg.merge_outputs(output_file, output_std)
+                .run_async(pipe_stdout=True, pipe_stderr=True)
+            )
+
+        return process
+
+    def save_chunk(self, n_audio_file):
+        """
+        Saves the current audio frames to a WAV file in a separate thread.
+
+        Args:
+        n_audio_file (int): The index of the audio file which determines the filename.
+                            This helps in maintaining the order and uniqueness of each chunk.
+        """
+        t = threading.Thread(
+            target=self.write_audio_frames_to_file,
+            args=(self.frames[:], f"chunks/{n_audio_file}.wav",),
+        )
+        t.start()
+
+    def finalize_recording(self, n_audio_file):
+        """
+        Finalizes the recording process by saving any remaining audio frames,
+        closing the audio stream, and terminating the process.
+
+        Args:
+        n_audio_file (int): The file index to be used if there are remaining audio frames to be saved.
+                            This index is incremented before use if the last chunk is saved.
+        """
+        if self.save_output_recording and len(self.frames):
+            self.write_audio_frames_to_file(
+                self.frames[:], f"chunks/{n_audio_file}.wav"
+            )
+            n_audio_file += 1
+        self.stream.stop_stream()
+        self.stream.close()
+        self.p.terminate()
+        self.close_all_clients()
+        if self.save_output_recording:
+            self.write_output_recording(n_audio_file)
+        self.write_all_clients_srt()
+
+    def record(self):
         """
         Record audio data from the input stream and save it to a WAV file.
 
@@ -445,15 +525,12 @@ class TranscriptionTeeClient:
         The recording will continue until the specified duration is reached or until the `RECORDING` flag is set to `False`.
         The recording process can be interrupted by sending a KeyboardInterrupt (e.g., pressing Ctrl+C). After recording,
         the method combines all the saved audio chunks into the specified `out_file`.
-
-        Args:
-            out_file (str, optional): The name of the output WAV file to save the entire recording.
-                                      Default is "output_recording.wav".
-
         """
         n_audio_file = 0
-        if not os.path.exists("chunks"):
-            os.makedirs("chunks", exist_ok=True)
+        if self.save_output_recording:
+            if os.path.exists("chunks"):
+                shutil.rmtree("chunks")
+            os.makedirs("chunks")
         try:
             for _ in range(0, int(self.rate / self.chunk * self.record_seconds)):
                 if not any(client.recording for client in self.clients):
@@ -467,31 +544,14 @@ class TranscriptionTeeClient:
 
                 # save frames if more than a minute
                 if len(self.frames) > 60 * self.rate:
-                    t = threading.Thread(
-                        target=self.write_audio_frames_to_file,
-                        args=(
-                            self.frames[:],
-                            f"chunks/{n_audio_file}.wav",
-                        ),
-                    )
-                    t.start()
-                    n_audio_file += 1
+                    if self.save_output_recording:
+                        self.save_chunk(n_audio_file)
+                        n_audio_file += 1
                     self.frames = b""
             self.write_all_clients_srt()
 
         except KeyboardInterrupt:
-            if len(self.frames):
-                self.write_audio_frames_to_file(
-                    self.frames[:], f"chunks/{n_audio_file}.wav"
-                )
-                n_audio_file += 1
-            self.stream.stop_stream()
-            self.stream.close()
-            self.p.terminate()
-            self.close_all_clients()
-
-            self.write_output_recording(n_audio_file, out_file)
-            self.write_all_clients_srt()
+            self.finalize_recording(n_audio_file)
 
     def write_audio_frames_to_file(self, frames, file_name):
         """
@@ -512,7 +572,7 @@ class TranscriptionTeeClient:
             wavfile.setframerate(self.rate)
             wavfile.writeframes(frames)
 
-    def write_output_recording(self, n_audio_file, out_file):
+    def write_output_recording(self, n_audio_file):
         """
         Combine and save recorded audio chunks into a single WAV file.
 
@@ -531,7 +591,7 @@ class TranscriptionTeeClient:
             for i in range(n_audio_file)
             if os.path.exists(f"chunks/{i}.wav")
         ]
-        with wave.open(out_file, "wb") as wavfile:
+        with wave.open(self.output_recording_filename, "wb") as wavfile:
             wavfile: wave.Wave_write
             wavfile.setnchannels(self.channels)
             wavfile.setsampwidth(2)
@@ -546,6 +606,9 @@ class TranscriptionTeeClient:
                 # remove this file
                 os.remove(in_file)
         wavfile.close()
+        # clean up temporary directory to store chunks
+        if os.path.exists("chunks"):
+            shutil.rmtree("chunks")
 
     @staticmethod
     def bytes_to_float_array(audio_bytes):
@@ -564,6 +627,7 @@ class TranscriptionTeeClient:
         raw_data = np.frombuffer(buffer=audio_bytes, dtype=np.int16)
         return raw_data.astype(np.float32) / 32768.0
 
+
 class TranscriptionClient(TranscriptionTeeClient):
     """
     Client for handling audio transcription tasks via a single WebSocket connection.
@@ -576,6 +640,9 @@ class TranscriptionClient(TranscriptionTeeClient):
         port (int): The port number to connect to on the server.
         lang (str, optional): The primary language for transcription. Default is None, which defaults to English ('en').
         translate (bool, optional): Indicates whether translation tasks are required (default is False).
+        save_output_recording (bool, optional): Indicates whether to save recording from microphone.
+        output_recording_filename (str, optional): File to save the output recording.
+        output_transcription_path (str, optional): File to save the output transcription.
 
     Attributes:
         client (Client): An instance of the underlying Client class responsible for handling the WebSocket connection.
@@ -587,6 +654,26 @@ class TranscriptionClient(TranscriptionTeeClient):
         transcription_client()
         ```
     """
-    def __init__(self, host, port, lang=None, translate=False, model="small", use_vad=True):
-        self.client = Client(host, port, lang, translate, model, srt_file_path="output.srt", use_vad=use_vad)
-        TranscriptionTeeClient.__init__(self, [self.client])
+    def __init__(
+        self,
+        host,
+        port,
+        lang=None,
+        translate=False,
+        model="small",
+        use_vad=True,
+        save_output_recording=False,
+        output_recording_filename="./output_recording.wav",
+        output_transcription_path="./output.srt"
+    ):
+        self.client = Client(host, port, lang, translate, model, srt_file_path=output_transcription_path, use_vad=use_vad)
+        if save_output_recording and not output_recording_filename.endswith(".wav"):
+            raise ValueError(f"Please provide a valid `output_recording_filename`: {output_recording_filename}")
+        if not output_transcription_path.endswith(".srt"):
+            raise ValueError(f"Please provide a valid `output_transcription_path`: {output_transcription_path}. The file extension should be `.srt`.")
+        TranscriptionTeeClient.__init__(
+            self,
+            [self.client],
+            save_output_recording=save_output_recording,
+            output_recording_filename=output_recording_filename
+        )

whisper_live/server.py

@@ -4,6 +4,9 @@ import threading
 import json
 import functools
 import logging
+from enum import Enum
+from typing import List, Optional
+
 import torch
 import numpy as np
 from websockets.sync.server import serve
@@ -121,6 +124,25 @@ class ClientManager:
         return False
 
 
+class BackendType(Enum):
+    FASTER_WHISPER = "faster_whisper"
+    TENSORRT = "tensorrt"
+
+    @staticmethod
+    def valid_types() -> List[str]:
+        return [backend_type.value for backend_type in BackendType]
+
+    @staticmethod
+    def is_valid(backend: str) -> bool:
+        return backend in BackendType.valid_types()
+
+    def is_faster_whisper(self) -> bool:
+        return self == BackendType.FASTER_WHISPER
+
+    def is_tensorrt(self) -> bool:
+        return self == BackendType.TENSORRT
+
+
 class TranscriptionServer:
     RATE = 16000
 
@@ -128,12 +150,15 @@ class TranscriptionServer:
         self.client_manager = ClientManager()
         self.no_voice_activity_chunks = 0
         self.use_vad = True
+        self.single_model = False
 
     def initialize_client(
         self, websocket, options, faster_whisper_custom_model_path,
         whisper_tensorrt_path, trt_multilingual
     ):
-        if self.backend == "tensorrt":
+        client: Optional[ServeClientBase] = None
+
+        if self.backend.is_tensorrt():
             try:
                 client = ServeClientTensorRT(
                     websocket,
@@ -141,7 +166,8 @@ class TranscriptionServer:
                     language=options["language"],
                     task=options["task"],
                     client_uid=options["uid"],
-                    model=whisper_tensorrt_path
+                    model=whisper_tensorrt_path,
+                    single_model=self.single_model,
                 )
                 logging.info("Running TensorRT backend.")
             except Exception as e:
@@ -153,9 +179,9 @@ class TranscriptionServer:
                     "message": "TensorRT-LLM not supported on Server yet. "
                                "Reverting to available backend: 'faster_whisper'"
                 }))
-                self.backend = "faster_whisper"
+                self.backend = BackendType.FASTER_WHISPER
 
-        if self.backend == "faster_whisper":
+        if self.backend.is_faster_whisper():
             if faster_whisper_custom_model_path is not None and os.path.exists(faster_whisper_custom_model_path):
                 logging.info(f"Using custom model {faster_whisper_custom_model_path}")
                 options["model"] = faster_whisper_custom_model_path
@@ -168,9 +194,13 @@ class TranscriptionServer:
                 initial_prompt=options.get("initial_prompt"),
                 vad_parameters=options.get("vad_parameters"),
                 use_vad=self.use_vad,
+                single_model=self.single_model,
             )
             logging.info("Running faster_whisper backend.")
 
+        if client is None:
+            raise ValueError(f"Backend type {self.backend.value} not recognised or not handled.")
+
         self.client_manager.add_client(websocket, client)
 
     def get_audio_from_websocket(self, websocket):
@@ -199,8 +229,7 @@ class TranscriptionServer:
                 websocket.close()
                 return False  # Indicates that the connection should not continue
 
-            if self.backend == "tensorrt":
-                self.vad_detector = VoiceActivityDetector(frame_rate=self.RATE)
+            self.vad_detector = VoiceActivityDetector(frame_rate=self.RATE)
             self.initialize_client(websocket, options, faster_whisper_custom_model_path,
                                    whisper_tensorrt_path, trt_multilingual)
             return True
@@ -218,24 +247,22 @@ class TranscriptionServer:
         frame_np = self.get_audio_from_websocket(websocket)
         client = self.client_manager.get_client(websocket)
         if frame_np is False:
-            if self.backend == "tensorrt":
-                client.set_eos(True)
+            client.set_eos(True)
             return False
 
-        if self.backend == "tensorrt":
-            voice_active = self.voice_activity(websocket, frame_np)
-            if voice_active:
-                self.no_voice_activity_chunks = 0
-                client.set_eos(False)
-            if self.use_vad and not voice_active:
-                return True
+        voice_active = self.voice_activity(websocket, frame_np)
+        if voice_active:
+            self.no_voice_activity_chunks = 0
+            client.set_eos(False)
+        if self.use_vad and not voice_active:
+            return True
 
         client.add_frames(frame_np)
         return True
 
     def recv_audio(self,
                    websocket,
-                   backend="faster_whisper",
+                   backend: BackendType = BackendType.FASTER_WHISPER,
                    faster_whisper_custom_model_path=None,
                    whisper_tensorrt_path=None,
                    trt_multilingual=False):
@@ -288,7 +315,8 @@ class TranscriptionServer:
             backend="tensorrt",
             faster_whisper_custom_model_path=None,
             whisper_tensorrt_path=None,
-            trt_multilingual=False):
+            trt_multilingual=False,
+            single_model=False):
         """
         Run the transcription server.
 
@@ -296,10 +324,18 @@ class TranscriptionServer:
             host (str): The host address to bind the server.
             port (int): The port number to bind the server.
         """
+        if faster_whisper_custom_model_path is not None and not os.path.exists(faster_whisper_custom_model_path):
+            raise ValueError(f"Custom faster_whisper model '{faster_whisper_custom_model_path}' is not a valid path.")
+        if whisper_tensorrt_path is not None and not os.path.exists(whisper_tensorrt_path):
+            raise ValueError(f"TensorRT model '{whisper_tensorrt_path}' is not a valid path.")
+
+        self.single_model = single_model
+        if not BackendType.is_valid(backend):
+            raise ValueError(f"{backend} is not a valid backend type. Choose backend from {BackendType.valid_types()}")
         with serve(
             functools.partial(
                 self.recv_audio,
-                backend=backend,
+                backend=BackendType(backend),
                 faster_whisper_custom_model_path=faster_whisper_custom_model_path,
                 whisper_tensorrt_path=whisper_tensorrt_path,
                 trt_multilingual=trt_multilingual
@@ -372,6 +408,7 @@ class ServeClientBase(object):
         self.add_pause_thresh = 3       # add a blank to segment list as a pause(no speech) for 3 seconds
         self.transcript = []
         self.send_last_n_segments = 10
+        self.eos = False
 
         # text formatting
         self.pick_previous_segments = 2
@@ -379,6 +416,18 @@ class ServeClientBase(object):
         # threading
         self.lock = threading.Lock()
 
+    def set_eos(self, eos):
+        """
+        Sets the End of Speech (EOS) flag.
+
+        Args:
+            eos (bool): The value to set for the EOS flag.
+        """
+        self.lock.acquire()
+        self.eos = eos
+        self.lock.release()
+
+
     def speech_to_text(self):
         raise NotImplementedError
 
@@ -499,7 +548,8 @@ class ServeClientBase(object):
             self.websocket.send(
                 json.dumps({
                     "uid": self.client_uid,
-                    "segments": segments,
+                    "text": segments,
+                    "eos": self.eos
                 })
             )
         except Exception as e:
@@ -532,7 +582,11 @@ class ServeClientBase(object):
 
 
 class ServeClientTensorRT(ServeClientBase):
-    def __init__(self, websocket, task="transcribe", multilingual=False, language=None, client_uid=None, model=None):
+
+    SINGLE_MODEL = None
+    SINGLE_MODEL_LOCK = threading.Lock()
+
+    def __init__(self, websocket, task="transcribe", multilingual=False, language=None, client_uid=None, model=None, single_model=False):
         """
         Initialize a ServeClient instance.
         The Whisper model is initialized based on the client's language and device availability.
@@ -546,21 +600,22 @@ class ServeClientTensorRT(ServeClientBase):
             multilingual (bool, optional): Whether the client supports multilingual transcription. Defaults to False.
             language (str, optional): The language for transcription. Defaults to None.
             client_uid (str, optional): A unique identifier for the client. Defaults to None.
+            single_model (bool, optional): Whether to instantiate a new model for each client connection. Defaults to False.
 
         """
         super().__init__(client_uid, websocket)
         self.language = language if multilingual else "en"
         self.task = task
         self.eos = False
-        self.transcriber = WhisperTRTLLM(
-            model,
-            assets_dir="assets",
-            device="cuda",
-            is_multilingual=multilingual,
-            language=self.language,
-            task=self.task
-        )
-        self.warmup()
+
+        if single_model:
+            if ServeClientTensorRT.SINGLE_MODEL is None:
+                self.create_model(model, multilingual)
+                ServeClientTensorRT.SINGLE_MODEL = self.transcriber
+            else:
+                self.transcriber = ServeClientTensorRT.SINGLE_MODEL
+        else:
+            self.create_model(model, multilingual)
 
         # threading
         self.trans_thread = threading.Thread(target=self.speech_to_text)
@@ -572,6 +627,21 @@ class ServeClientTensorRT(ServeClientBase):
             "backend": "tensorrt"
         }))
 
+    def create_model(self, model, multilingual, warmup=True):
+        """
+        Instantiates a new model, sets it as the transcriber and does warmup if desired.
+        """
+        self.transcriber = WhisperTRTLLM(
+            model,
+            assets_dir="assets",
+            device="cuda",
+            is_multilingual=multilingual,
+            language=self.language,
+            task=self.task
+        )
+        if warmup:
+            self.warmup()
+
     def warmup(self, warmup_steps=10):
         """
         Warmup TensorRT since first few inferences are slow.
@@ -584,17 +654,6 @@ class ServeClientTensorRT(ServeClientBase):
         for i in range(warmup_steps):
             self.transcriber.transcribe(mel)
 
-    def set_eos(self, eos):
-        """
-        Sets the End of Speech (EOS) flag.
-
-        Args:
-            eos (bool): The value to set for the EOS flag.
-        """
-        self.lock.acquire()
-        self.eos = eos
-        self.lock.release()
-
     def handle_transcription_output(self, last_segment, duration):
         """
         Handle the transcription output, updating the transcript and sending data to the client.
@@ -616,12 +675,16 @@ class ServeClientTensorRT(ServeClientBase):
         Args:
             input_bytes (np.array): The audio chunk to transcribe.
         """
+        if ServeClientTensorRT.SINGLE_MODEL:
+            ServeClientTensorRT.SINGLE_MODEL_LOCK.acquire()
         logging.info(f"[WhisperTensorRT:] Processing audio with duration: {input_bytes.shape[0] / self.RATE}")
         mel, duration = self.transcriber.log_mel_spectrogram(input_bytes)
         last_segment = self.transcriber.transcribe(
             mel,
             text_prefix=f"<|startoftranscript|><|{self.language}|><|{self.task}|><|notimestamps|>"
         )
+        if ServeClientTensorRT.SINGLE_MODEL:
+            ServeClientTensorRT.SINGLE_MODEL_LOCK.release()
         if last_segment:
             self.handle_transcription_output(last_segment, duration)
 
@@ -681,8 +744,12 @@ class ServeClientTensorRT(ServeClientBase):
 
 
 class ServeClientFasterWhisper(ServeClientBase):
+
+    SINGLE_MODEL = None
+    SINGLE_MODEL_LOCK = threading.Lock()
+
     def __init__(self, websocket, task="transcribe", device=None, language=None, client_uid=None, model="small.en",
-                 initial_prompt=None, vad_parameters=None, use_vad=True):
+                 initial_prompt=None, vad_parameters=None, use_vad=True, single_model=False):
         """
         Initialize a ServeClient instance.
         The Whisper model is initialized based on the client's language and device availability.
@@ -697,6 +764,7 @@ class ServeClientFasterWhisper(ServeClientBase):
             client_uid (str, optional): A unique identifier for the client. Defaults to None.
             model (str, optional): The whisper model size. Defaults to 'small.en'
             initial_prompt (str, optional): Prompt for whisper inference. Defaults to None.
+            single_model (bool, optional): Whether to instantiate a new model for each client connection. Defaults to False.
         """
         super().__init__(client_uid, websocket)
         self.model_sizes = [
@@ -711,19 +779,23 @@ class ServeClientFasterWhisper(ServeClientBase):
         self.task = task
         self.initial_prompt = initial_prompt
         self.vad_parameters = vad_parameters or {"threshold": 0.5}
-        self.no_speech_thresh = 0.45
+        self.no_speech_thresh = 0.35
 
         device = "cuda" if torch.cuda.is_available() else "cpu"
 
         if self.model_size_or_path is None:
             return
 
-        self.transcriber = WhisperModel(
-            self.model_size_or_path,
-            device=device,
-            compute_type="int8" if device == "cpu" else "float16",
-            local_files_only=False,
-        )
+        if single_model:
+            if ServeClientFasterWhisper.SINGLE_MODEL is None:
+                self.create_model(device)
+                ServeClientFasterWhisper.SINGLE_MODEL = self.transcriber
+            else:
+                print("Re-using already initialized model.")
+                self.transcriber = ServeClientFasterWhisper.SINGLE_MODEL
+        else:
+            self.create_model(device)
+
         self.use_vad = use_vad
 
         # threading
@@ -739,6 +811,17 @@ class ServeClientFasterWhisper(ServeClientBase):
             )
         )
 
+    def create_model(self, device):
+        """
+        Instantiates a new model, sets it as the transcriber.
+        """
+        self.transcriber = WhisperModel(
+            self.model_size_or_path,
+            device=device,
+            compute_type="int8" if device == "cpu" else "float16",
+            local_files_only=False,
+        )
+
     def check_valid_model(self, model_size):
         """
         Check if it's a valid whisper model size.
@@ -794,13 +877,18 @@ class ServeClientFasterWhisper(ServeClientBase):
             depends on the implementation of the `transcriber.transcribe` method but typically
             includes the transcribed text.
         """
+        if ServeClientFasterWhisper.SINGLE_MODEL:
+            ServeClientFasterWhisper.SINGLE_MODEL_LOCK.acquire()
         result, info = self.transcriber.transcribe(
             input_sample,
             initial_prompt=self.initial_prompt,
             language=self.language,
             task=self.task,
-            vad_filter=self.use_vad,
-            vad_parameters=self.vad_parameters if self.use_vad else None)
+            vad_filter=False,
+            vad_parameters=self.vad_parameters if self.use_vad else None,
+            beam_size=5)
+        if ServeClientFasterWhisper.SINGLE_MODEL:
+            ServeClientFasterWhisper.SINGLE_MODEL_LOCK.release()
 
         if self.language is None and info is not None:
             self.set_language(info)
@@ -841,17 +929,16 @@ class ServeClientFasterWhisper(ServeClientBase):
             result (str): The result from whisper inference i.e. the list of segments.
             duration (float): Duration of the transcribed audio chunk.
         """
-        segments = []
         if len(result):
-            self.t_start = None
             last_segment = self.update_segments(result, duration)
-            segments = self.prepare_segments(last_segment)
-        else:
-            # show previous output if there is pause i.e. no output from whisper
-            segments = self.get_previous_output()
 
-        if len(segments):
-            self.send_transcription_to_client(segments)
+        if len(self.text):
+            if self.eos and last_segment is None:
+                self.send_transcription_to_client(' '.join([s.strip() for s in self.text]))
+                self.set_eos(False)
+                self.text = []
+            elif not self.eos:
+                self.send_transcription_to_client(' '.join([s.strip() for s in self.text]))
 
     def speech_to_text(self):
         """
@@ -881,7 +968,12 @@ class ServeClientFasterWhisper(ServeClientBase):
             self.clip_audio_if_no_valid_segment()
 
             input_bytes, duration = self.get_audio_chunk_for_processing()
-            if duration < 1.0:
+            if duration < 0.6:
+                if len(self.text) and self.eos:
+                    self.send_transcription_to_client(' '.join([s.strip() for s in self.text]))
+                    self.set_eos(False)
+                    self.text = []
+                time.sleep(0.1)
                 continue
             try:
                 input_sample = input_bytes.copy()
@@ -889,7 +981,7 @@ class ServeClientFasterWhisper(ServeClientBase):
 
                 if result is None or self.language is None:
                     self.timestamp_offset += duration
-                    time.sleep(0.25)    # wait for voice activity, result is None when no voice activity
+                    time.sleep(0.1)    # wait for voice activity, result is None when no voice activity
                     continue
                 self.handle_transcription_output(result, duration)
 
@@ -938,13 +1030,13 @@ class ServeClientFasterWhisper(ServeClientBase):
             dict or None: The last processed segment with its start time, end time, and transcribed text.
                      Returns None if there are no valid segments to process.
         """
+        last_segment = None
         offset = None
         self.current_out = ''
         # process complete segments
         if len(segments) > 1:
             for i, s in enumerate(segments[:-1]):
                 text_ = s.text
-                self.text.append(text_)
                 start, end = self.timestamp_offset + s.start, self.timestamp_offset + min(duration, s.end)
 
                 if start >= end:
@@ -952,15 +1044,17 @@ class ServeClientFasterWhisper(ServeClientBase):
                 if s.no_speech_prob > self.no_speech_thresh:
                     continue
 
+                self.text.append(text_)
                 self.transcript.append(self.format_segment(start, end, text_))
                 offset = min(duration, s.end)
 
-        self.current_out += segments[-1].text
-        last_segment = self.format_segment(
-            self.timestamp_offset + segments[-1].start,
-            self.timestamp_offset + min(duration, segments[-1].end),
-            self.current_out
-        )
+        if segments[-1].no_speech_prob <= self.no_speech_thresh:
+            self.current_out += segments[-1].text
+            last_segment = self.format_segment(
+                self.timestamp_offset + segments[-1].start,
+                self.timestamp_offset + min(duration, segments[-1].end),
+                self.current_out
+            )
 
         # if same incomplete segment is seen multiple times then update the offset
         # and append the segment to the list
@@ -969,7 +1063,7 @@ class ServeClientFasterWhisper(ServeClientBase):
         else:
             self.same_output_threshold = 0
 
-        if self.same_output_threshold > 5:
+        if self.same_output_threshold > 2:
             if not len(self.text) or self.text[-1].strip().lower() != self.current_out.strip().lower():
                 self.text.append(self.current_out)
                 self.transcript.append(self.format_segment(

whisper_live/transcriber.py

@@ -496,7 +496,7 @@ class WhisperModel:
                 content_frames - seek,
                 seek_clip_end - seek,
             )
-            segment = features[:, seek : seek + segment_size]
+            segment = features[:, seek:seek + segment_size]
             segment_duration = segment_size * self.feature_extractor.time_per_frame
             segment = pad_or_trim(segment, self.feature_extractor.nb_max_frames)
 
@@ -685,7 +685,7 @@ class WhisperModel:
                             continue
                         if is_segment_anomaly(segment):
                             next_segment = next_words_segment(
-                                current_segments[si + 1 :]
+                                current_segments[si + 1:]
                             )
                             if next_segment is not None:
                                 hal_next_start = next_segment["words"][0]["start"]
@@ -909,7 +909,7 @@ class WhisperModel:
 
         if previous_tokens:
             prompt.append(tokenizer.sot_prev)
-            prompt.extend(previous_tokens[-(self.max_length // 2 - 1) :])
+            prompt.extend(previous_tokens[-(self.max_length // 2 - 1):])
 
         prompt.extend(tokenizer.sot_sequence)
 
@@ -926,7 +926,7 @@ class WhisperModel:
 
         return prompt
 
-    def add_word_timestamps(                                                    # noqa: C901
+    def add_word_timestamps(
         self,
         segments: List[dict],
         tokenizer: Tokenizer,

whisper_live/transcriber_tensorrt.py

@@ -35,6 +35,8 @@ class WhisperEncoding:
         with open(config_path, 'r') as f:
             config = json.load(f)
 
+        use_gpt_attention_plugin = config['plugin_config'][
+            'gpt_attention_plugin']
         dtype = config['builder_config']['precision']
         n_mels = config['builder_config']['n_mels']
         num_languages = config['builder_config']['num_languages']
@@ -51,12 +53,21 @@ class WhisperEncoding:
         return session
 
     def get_audio_features(self, mel):
-        inputs = OrderedDict()
-        output_list = []
 
-        inputs.update({'x': mel})
-        output_list.append(
-            TensorInfo('x', str_dtype_to_trt(self.dtype), mel.shape))
+        input_lengths = torch.tensor(
+            [mel.shape[2] // 2 for _ in range(mel.shape[0])],
+            dtype=torch.int32,
+            device=mel.device)
+
+        inputs = OrderedDict()
+        inputs['x'] = mel
+        inputs['input_lengths'] = input_lengths
+
+        output_list = [
+            TensorInfo('x', str_dtype_to_trt(self.dtype), mel.shape),
+            TensorInfo('input_lengths', str_dtype_to_trt('int32'),
+                       input_lengths.shape)
+        ]
 
         output_info = (self.session).infer_shapes(output_list)
 
@@ -101,6 +112,8 @@ class WhisperDecoding:
             decoder_engine_buffer = f.read()
 
         decoder_model_config = ModelConfig(
+            max_batch_size=self.decoder_config['max_batch_size'],
+            max_beam_width=self.decoder_config['max_beam_width'],
             num_heads=self.decoder_config['num_heads'],
             num_kv_heads=self.decoder_config['num_heads'],
             hidden_size=self.decoder_config['hidden_size'],
@@ -141,6 +154,10 @@ class WhisperDecoding:
                                              device='cuda')
         decoder_max_input_length = torch.max(decoder_input_lengths).item()
 
+        cross_attention_mask = torch.ones(
+            [encoder_outputs.shape[0], 1,
+             encoder_outputs.shape[1]]).int().cuda()
+
         # generation config
         sampling_config = SamplingConfig(end_id=eot_id,
                                          pad_id=eot_id,
@@ -161,6 +178,7 @@ class WhisperDecoding:
             sampling_config,
             encoder_output=encoder_outputs,
             encoder_input_lengths=encoder_input_lengths,
+            cross_attention_mask=cross_attention_mask,
         )
         torch.cuda.synchronize()
 

whisper_live/vad.py

@@ -94,7 +94,7 @@ class VoiceActivityDetection():
         return stacked.cpu()
 
     @staticmethod
-    def download(model_url="https://github.com/snakers4/silero-vad/raw/master/files/silero_vad.onnx"):
+    def download(model_url="https://github.com/snakers4/silero-vad/raw/v4.0/files/silero_vad.onnx"):
         target_dir = os.path.expanduser("~/.cache/whisper-live/")
 
         # Ensure the target directory exists