Bump version v0.6.1

Fix package name

.github/workflows/ci.yml

@@ -101,7 +101,7 @@ jobs:
   
   build-and-push-docker-tensorrt:
     needs: [run-tests, check-code-format]
-    timeout-minutes: 20
+    timeout-minutes: 60
     runs-on: ubuntu-22.04
     if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/'))
     steps:

README.md

@@ -77,6 +77,9 @@ If you don't want this, set `--no_single_model`.
   - `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`.
+  - `max_clients`: Specifies the maximum number of clients the server should allow. Defaults to 4.
+  - `max_connection_time`: Maximum connection time for each client in seconds. Defaults to 600.
+
 ```python
 from whisper_live.client import TranscriptionClient
 client = TranscriptionClient(
@@ -84,10 +87,12 @@ client = TranscriptionClient(
   9090,
   lang="en",
   translate=False,
-  model="small",
+  model="small",                                      # also support hf_model => `Systran/faster-whisper-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
+  output_recording_filename="./output_recording.wav", # Only used for microphone input
+  max_clients=4,
+  max_connection_time=600
 )
 ```
 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.
@@ -127,13 +132,17 @@ client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/b
   ```bash
   docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt
 
-  # Build tiny.en engine
+  # Build small.en engine
-  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en
+  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en        # float16
+  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int8   # int8 weight only quantization
+  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int4   # int4 weight only quantization
 
-  # Run server with tiny.en
+  # Run server with small.en
   python3 run_server.py --port 9090 \
                         --backend tensorrt \
-                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en"
+                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_float16"
+                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int8"
+                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int4"
   ```
 
 - CPU

TensorRT_whisper.md

@@ -1,17 +1,11 @@
 # WhisperLive-TensorRT
 We have only tested the TensorRT backend in docker so, we recommend docker for a smooth TensorRT backend setup.
-**Note**: We use `tensorrt_llm==0.9.0`
+**Note**: We use `tensorrt_llm==0.15.0.dev2024111200`
 
 ## 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)
 
-- Clone this repo.
-```bash
-git clone https://github.com/collabora/WhisperLive.git
-cd WhisperLive
-```
-
 - Run WhisperLive TensorRT in docker
 ```bash
 docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt:latest
@@ -21,7 +15,9 @@ docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it g
 - 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 build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en
+bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en        # float16
+bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int8   # int8 weight only quantization
+bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int4   # int4 weight only quantization
 
 # convert small multilingual model
 bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small
@@ -32,11 +28,11 @@ bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small
 # Run English only model
 python3 run_server.py --port 9090 \
                       --backend tensorrt \
-                      --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en"
+                      --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_float16"
 
 # Run Multilingual model
 python3 run_server.py --port 9090 \
                       --backend tensorrt \
-                      --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small" \
+                      --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_float16" \
                       --trt_multilingual
 ```

docker/Dockerfile.tensorrt

@@ -1,19 +1,22 @@
-FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
+FROM nvidia/cuda:12.5.1-runtime-ubuntu22.04 AS base
 
 ARG DEBIAN_FRONTEND=noninteractive
 
 RUN apt-get update && apt-get install -y \
-    python3.10 python3-pip openmpi-bin libopenmpi-dev git wget \
+    python3.10 python3-pip openmpi-bin libopenmpi-dev git git-lfs 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
+FROM base AS devel
-
+RUN pip3 install --no-cache-dir -U tensorrt_llm==0.15.0.dev2024111200 --extra-index-url https://pypi.nvidia.com
 WORKDIR /app
-
+RUN git clone https://github.com/NVIDIA/TensorRT-LLM.git && cd TensorRT-LLM && \
-RUN git clone -b v0.9.0 --depth 1 https://github.com/NVIDIA/TensorRT-LLM.git && \
+    git checkout c629546ce429623c8a163633095230154a6f0574 && cd ../ && \
     mv TensorRT-LLM/examples ./TensorRT-LLM-examples && \
     rm -rf TensorRT-LLM
 
+
+FROM devel AS release
+WORKDIR /app
 COPY assets/ ./assets
 RUN wget -nc -P assets/ https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz
 
@@ -22,7 +25,6 @@ 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

@@ -1,9 +1,7 @@
-faster-whisper==1.0.1
+faster-whisper==1.1.0
-torch
 websockets
 onnxruntime==1.16.0
 numba
-openai-whisper
 kaldialign
 soundfile
 ffmpeg-python
@@ -11,3 +9,5 @@ scipy
 jiwer
 evaluate
 numpy<2
+openai-whisper==20240930
+tokenizers==0.20.3

scripts/build_whisper_tensorrt.sh

@@ -38,12 +38,24 @@ download_and_build_model() {
         "large-v3" | "large")
             model_url="https://openaipublic.azureedge.net/main/whisper/models/e5b1a55b89c1367dacf97e3e19bfd829a01529dbfdeefa8caeb59b3f1b81dadb/large-v3.pt"
             ;;
+        "large-v3-turbo" | "turbo")
+            model_url="https://openaipublic.azureedge.net/main/whisper/models/aff26ae408abcba5fbf8813c21e62b0941638c5f6eebfb145be0c9839262a19a/large-v3-turbo.pt"
+            ;;
         *)
             echo "Invalid model name: $model_name"
             exit 1
             ;;
     esac
 
+    if [ "$model_name" == "turbo" ]; then
+        model_name="large-v3-turbo"
+    fi
+
+    local inference_precision="float16"
+    local weight_only_precision="${2:-float16}"
+    local max_beam_width=4
+    local max_batch_size=1
+
     echo "Downloading $model_name..."
     # wget --directory-prefix=assets "$model_url"
     # echo "Download completed: ${model_name}.pt"
@@ -54,11 +66,43 @@ download_and_build_model() {
         echo "${model_name}.pt already exists in assets directory."
     fi
 
-    local output_dir="whisper_${model_name//./_}"
+    local sanitized_model_name="${model_name//./_}"
+    local checkpoint_dir="whisper_${sanitized_model_name}_weights_${weight_only_precision}"
+    local output_dir="whisper_${sanitized_model_name}_${weight_only_precision}"
     echo "$output_dir"
-    echo "Running build script for $model_name with output directory $output_dir"
+    echo "Converting model weights for $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"
+    python3 convert_checkpoint.py \
-    echo "Whisper $model_name TensorRT engine built."
+        $( [[ "$weight_only_precision" == "int8" || "$weight_only_precision" == "int4" ]] && echo "--use_weight_only --weight_only_precision $weight_only_precision" ) \
+        --output_dir "$checkpoint_dir" --model_name "$model_name"
+    
+    echo "Building encoder for $model_name..."
+    trtllm-build \
+        --checkpoint_dir "${checkpoint_dir}/encoder" \
+        --output_dir "${output_dir}/encoder" \
+        --moe_plugin disable \
+        --enable_xqa disable \
+        --max_batch_size "$max_batch_size" \
+        --gemm_plugin disable \
+        --bert_attention_plugin "$inference_precision" \
+        --max_input_len 3000 \
+        --max_seq_len 3000
+    
+    echo "Building decoder for $model_name..."
+    trtllm-build \
+        --checkpoint_dir "${checkpoint_dir}/decoder" \
+        --output_dir "${output_dir}/decoder" \
+        --moe_plugin disable \
+        --enable_xqa disable \
+        --max_beam_width "$max_beam_width" \
+        --max_batch_size "$max_batch_size" \
+        --max_seq_len 200 \
+        --max_input_len 14 \
+        --max_encoder_input_len 3000 \
+        --gemm_plugin "$inference_precision" \
+        --bert_attention_plugin "$inference_precision" \
+        --gpt_attention_plugin "$inference_precision"
+
+    echo "TensorRT LLM engine built for $model_name."
     echo "========================================="
     echo "Model is located at: $(pwd)/$output_dir"
 }
@@ -70,8 +114,9 @@ fi
 
 tensorrt_examples_dir="$1"
 model_name="${2:-small.en}"
+weight_only_precision="${3:-float16}"  # Default to float16 if not provided
 
-cd $1/whisper
+cd $tensorrt_examples_dir/whisper
 pip install --no-deps -r requirements.txt
 
-download_and_build_model "$model_name"
+download_and_build_model "$model_name" "$weight_only_precision"

setup.py

@@ -11,7 +11,7 @@ README = (HERE / "README.md").read_text()
 
 # This call to setup() does all the work
 setup(
-    name="whisper-live",
+    name="whisper_live",
     version=__version__,
     description="A nearly-live implementation of OpenAI's Whisper.",
     long_description=README,
@@ -43,7 +43,7 @@ setup(
     ),
     install_requires=[
         "PyAudio",
-        "faster-whisper==1.0.1",
+        "faster-whisper==1.1.0",
         "torch",
         "torchaudio",
         "websockets",
@@ -52,9 +52,10 @@ setup(
         "scipy",
         "websocket-client",
         "numba",
-        "openai-whisper",
+        "openai-whisper==20240930",
         "kaldialign",
         "soundfile",
+        "tokenizers==0.20.3"
     ],
     python_requires=">=3.8"
 )

tests/test_client.py

@@ -48,7 +48,9 @@ class TestClientCallbacks(BaseTestCase):
             "language": self.client.language,
             "task": self.client.task,
             "model": self.client.model,
-            "use_vad": True
+            "use_vad": True,
+            "max_clients": 4,
+            "max_connection_time": 600,
         })
         self.client.on_open(self.mock_ws_app)
         self.mock_ws_app.send.assert_called_with(expected_message)
@@ -66,15 +68,15 @@ class TestClientCallbacks(BaseTestCase):
         message = json.dumps({
             "uid": self.client.uid,
             "segments": [
-                {"start": 0, "end": 1, "text": "Test transcript"},
+                {"start": 0, "end": 1, "text": "Test transcript", "completed": True},
-                {"start": 1, "end": 2, "text": "Test transcript 2"},
+                {"start": 1, "end": 2, "text": "Test transcript 2", "completed": True},
-                {"start": 2, "end": 3, "text": "Test transcript 3"}
+                {"start": 2, "end": 3, "text": "Test transcript 3", "completed": True}
             ]
         })
         self.client.on_message(self.mock_ws_app, message)
 
         # Assert that the transcript was updated correctly
-        self.assertEqual(len(self.client.transcript), 2)
+        self.assertEqual(len(self.client.transcript), 3)
         self.assertEqual(self.client.transcript[1]['text'], "Test transcript 2")
 
     def test_on_close(self):

tests/test_server.py

@@ -5,10 +5,10 @@ import unittest
 from unittest import mock
 
 import numpy as np
-import evaluate
+import jiwer
 
 from websockets.exceptions import ConnectionClosed
-from whisper_live.server import TranscriptionServer
+from whisper_live.server import TranscriptionServer, BackendType, ClientManager
 from whisper_live.client import Client, TranscriptionClient, TranscriptionTeeClient
 from whisper.normalizers import EnglishTextNormalizer
 
@@ -16,6 +16,7 @@ from whisper.normalizers import EnglishTextNormalizer
 class TestTranscriptionServerInitialization(unittest.TestCase):
     def test_initialization(self):
         server = TranscriptionServer()
+        server.client_manager = ClientManager(max_clients=4, max_connection_time=600)
         self.assertEqual(server.client_manager.max_clients, 4)
         self.assertEqual(server.client_manager.max_connection_time, 600)
         self.assertDictEqual(server.client_manager.clients, {})
@@ -25,6 +26,7 @@ class TestTranscriptionServerInitialization(unittest.TestCase):
 class TestGetWaitTime(unittest.TestCase):
     def setUp(self):
         self.server = TranscriptionServer()
+        self.server.client_manager = ClientManager(max_clients=4, max_connection_time=600)
         self.server.client_manager.start_times = {
             'client1': time.time() - 120,
             'client2': time.time() - 300
@@ -49,7 +51,7 @@ class TestServerConnection(unittest.TestCase):
             'task': 'transcribe',
             'model': 'tiny.en'
         })
-        self.server.recv_audio(mock_websocket, "faster_whisper")
+        self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
 
     @mock.patch('websockets.WebSocketCommonProtocol')
     def test_recv_audio_exception_handling(self, mock_websocket):
@@ -61,7 +63,7 @@ class TestServerConnection(unittest.TestCase):
         }),  np.array([1, 2, 3]).tobytes()]
 
         with self.assertLogs(level="ERROR"):
-            self.server.recv_audio(mock_websocket, "faster_whisper")
+            self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
 
         self.assertNotIn(mock_websocket, self.server.client_manager.clients)
 
@@ -82,7 +84,6 @@ class TestServerInferenceAccuracy(unittest.TestCase):
         cls.server_process.wait()
 
     def setUp(self):
-        self.metric = evaluate.load("wer")
         self.normalizer = EnglishTextNormalizer()
 
     def check_prediction(self, srt_path):
@@ -94,11 +95,8 @@ class TestServerInferenceAccuracy(unittest.TestCase):
         gt_normalized = self.normalizer(gt)
 
         # calculate WER
-        wer = self.metric.compute(
+        wer_score = jiwer.wer(gt_normalized, prediction_normalized)
-            predictions=[prediction_normalized],
+        self.assertLess(wer_score, 0.05)
-            references=[gt_normalized]
-        )
-        self.assertLess(wer, 0.05)
 
     def test_inference(self):
         client = TranscriptionClient(
@@ -124,10 +122,10 @@ class TestExceptionHandling(unittest.TestCase):
 
     @mock.patch('websockets.WebSocketCommonProtocol')
     def test_connection_closed_exception(self, mock_websocket):
-        mock_websocket.recv.side_effect = ConnectionClosed(1001, "testing connection closed")
+        mock_websocket.recv.side_effect = ConnectionClosed(1001, "testing connection closed", rcvd_then_sent=mock.Mock())
 
         with self.assertLogs(level="INFO") as log:
-            self.server.recv_audio(mock_websocket, "faster_whisper")
+            self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
             self.assertTrue(any("Connection closed by client" in message for message in log.output))
 
     @mock.patch('websockets.WebSocketCommonProtocol')
@@ -135,7 +133,7 @@ class TestExceptionHandling(unittest.TestCase):
         mock_websocket.recv.return_value = "invalid json"
 
         with self.assertLogs(level="ERROR") as log:
-            self.server.recv_audio(mock_websocket, "faster_whisper")
+            self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
             self.assertTrue(any("Failed to decode JSON from client" in message for message in log.output))
 
     @mock.patch('websockets.WebSocketCommonProtocol')
@@ -143,7 +141,7 @@ class TestExceptionHandling(unittest.TestCase):
         mock_websocket.recv.side_effect = RuntimeError("Unexpected error")
 
         with self.assertLogs(level="ERROR") as log:
-            self.server.recv_audio(mock_websocket, "faster_whisper")
+            self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
             for message in log.output:
                 print(message)
             print()

whisper_live/__version__.py

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

whisper_live/client.py

@@ -2,6 +2,7 @@ import os
 import shutil
 import wave
 
+import logging
 import numpy as np
 import pyaudio
 import threading
@@ -28,7 +29,10 @@ class Client:
         translate=False,
         model="small",
         srt_file_path="output.srt",
-        use_vad=True
+        use_vad=True,
+        log_transcription=True,
+        max_clients=4,
+        max_connection_time=600,
     ):
         """
         Initializes a Client instance for audio recording and streaming to a server.
@@ -56,11 +60,13 @@ class Client:
         self.use_vad = use_vad
         self.last_segment = None
         self.last_received_segment = None
+        self.log_transcription = log_transcription
+        self.max_clients = max_clients
+        self.max_connection_time = max_connection_time
 
         if translate:
             self.task = "translate"
 
-        self.timestamp_offset = 0.0
         self.audio_bytes = None
 
         if host is not None and port is not None:
@@ -106,9 +112,9 @@ class Client:
         for i, seg in enumerate(segments):
             if not text or text[-1] != seg["text"]:
                 text.append(seg["text"])
-                if i == len(segments) - 1:
+                if i == len(segments) - 1 and not seg.get("completed", False):
                     self.last_segment = seg
-                elif (self.server_backend == "faster_whisper" and
+                elif (self.server_backend == "faster_whisper" and seg.get("completed", False) and
                       (not self.transcript or
                         float(seg['start']) >= float(self.transcript[-1]['end']))):
                     self.transcript.append(seg)
@@ -117,6 +123,7 @@ class Client:
             self.last_response_received = time.time()
             self.last_received_segment = segments[-1]["text"]
 
+        if self.log_transcription:
             # Truncate to last 3 entries for brevity.
             text = text[-3:]
             utils.clear_screen()
@@ -196,7 +203,9 @@ class Client:
                     "language": self.language,
                     "task": self.task,
                     "model": self.model,
-                    "use_vad": self.use_vad
+                    "use_vad": self.use_vad,
+                    "max_clients": self.max_clients,
+                    "max_connection_time": self.max_connection_time,
                 }
             )
         )
@@ -250,7 +259,9 @@ class Client:
 
         """
         if self.server_backend == "faster_whisper":
-            if (self.last_segment):
+            if not self.transcript and self.last_segment is not None:
+                self.transcript.append(self.last_segment)
+            elif self.last_segment and self.transcript[-1]["text"] != self.last_segment["text"]:
                 self.transcript.append(self.last_segment)
             utils.create_srt_file(self.transcript, output_path)
 
@@ -431,6 +442,8 @@ class TranscriptionTeeClient:
 
     def handle_ffmpeg_process(self, process, stream_type):
         print(f"[INFO]: Connecting to {stream_type} stream...")
+        stderr_thread = threading.Thread(target=self.consume_stderr, args=(process,))
+        stderr_thread.start()
         try:
             # Process the stream
             while True:
@@ -477,6 +490,16 @@ class TranscriptionTeeClient:
 
         return process
 
+    def consume_stderr(self, process):
+        """
+        Consume and log the stderr output of a process in a separate thread.
+
+        Args:
+            process (subprocess.Popen): The process whose stderr output will be logged.
+        """
+        for line in iter(process.stderr.readline, b""):
+            logging.debug(f'[STDERR]: {line.decode()}')
+
     def save_chunk(self, n_audio_file):
         """
         Saves the current audio frames to a WAV file in a separate thread.
@@ -664,9 +687,17 @@ class TranscriptionClient(TranscriptionTeeClient):
         use_vad=True,
         save_output_recording=False,
         output_recording_filename="./output_recording.wav",
-        output_transcription_path="./output.srt"
+        output_transcription_path="./output.srt",
+        log_transcription=True,
+        max_clients=4,
+        max_connection_time=600,
     ):
-        self.client = Client(host, port, lang, translate, model, srt_file_path=output_transcription_path, use_vad=use_vad)
+        self.client = Client(
+            host, port, lang, translate, model, srt_file_path=output_transcription_path,
+            use_vad=use_vad, log_transcription=log_transcription, max_clients=max_clients,
+            max_connection_time=max_connection_time
+        )
+
         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"):

whisper_live/server.py

@@ -147,7 +147,7 @@ class TranscriptionServer:
     RATE = 16000
 
     def __init__(self):
-        self.client_manager = ClientManager()
+        self.client_manager = None
         self.no_voice_activity_chunks = 0
         self.use_vad = True
         self.single_model = False
@@ -181,6 +181,7 @@ class TranscriptionServer:
                 }))
                 self.backend = BackendType.FASTER_WHISPER
 
+        try:
             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}")
@@ -196,7 +197,10 @@ class TranscriptionServer:
                     use_vad=self.use_vad,
                     single_model=self.single_model,
                 )
+
                 logging.info("Running faster_whisper backend.")
+        except Exception as e:
+            return
 
         if client is None:
             raise ValueError(f"Backend type {self.backend.value} not recognised or not handled.")
@@ -224,6 +228,12 @@ class TranscriptionServer:
             logging.info("New client connected")
             options = websocket.recv()
             options = json.loads(options)
+
+            if self.client_manager is None:
+                max_clients = options.get('max_clients', 4)
+                max_connection_time = options.get('max_connection_time', 600)
+                self.client_manager = ClientManager(max_clients, max_connection_time)
+
             self.use_vad = options.get('use_vad')
             if self.client_manager.is_server_full(websocket, options):
                 websocket.close()
@@ -411,7 +421,7 @@ class ServeClientBase(object):
         self.prev_out = ''
         self.t_start = None
         self.exit = False
-        self.same_output_threshold = 0
+        self.same_output_count = 0
         self.show_prev_out_thresh = 5   # if pause(no output from whisper) show previous output for 5 seconds
         self.add_pause_thresh = 3       # add a blank to segment list as a pause(no speech) for 3 seconds
         self.transcript = []
@@ -469,6 +479,7 @@ class ServeClientBase(object):
         Clip audio if the current chunk exceeds 30 seconds, this basically implies that
         no valid segment for the last 30 seconds from whisper
         """
+        with self.lock:
             if self.frames_np[int((self.timestamp_offset - self.frames_offset)*self.RATE):].shape[0] > 25 * self.RATE:
                 duration = self.frames_np.shape[0] / self.RATE
                 self.timestamp_offset = self.frames_offset + duration - 5
@@ -487,6 +498,7 @@ class ServeClientBase(object):
                 - input_bytes (np.ndarray): The next chunk of audio data to be processed.
                 - duration (float): The duration of the audio chunk in seconds.
         """
+        with self.lock:
             samples_take = max(0, (self.timestamp_offset - self.frames_offset) * self.RATE)
             input_bytes = self.frames_np[int(samples_take):].copy()
         duration = input_bytes.shape[0] / self.RATE
@@ -705,6 +717,8 @@ class ServeClientTensorRT(ServeClientBase):
             self.transcript.append({"text": last_segment + " "})
         elif self.transcript[-1]["text"].strip() != last_segment:
             self.transcript.append({"text": last_segment + " "})
+        
+        with self.lock:
             self.timestamp_offset += duration
 
     def speech_to_text(self):
@@ -774,23 +788,32 @@ class ServeClientFasterWhisper(ServeClientBase):
         super().__init__(client_uid, websocket)
         self.model_sizes = [
             "tiny", "tiny.en", "base", "base.en", "small", "small.en",
-            "medium", "medium.en", "large-v2", "large-v3",
+            "medium", "medium.en", "large-v2", "large-v3", "distil-small.en",
+            "distil-medium.en", "distil-large-v2", "distil-large-v3",
+            "large-v3-turbo", "turbo"
         ]
-        if not os.path.exists(model):
+
-            self.model_size_or_path = self.check_valid_model(model)
-        else:
         self.model_size_or_path = model
         self.language = "en" if self.model_size_or_path.endswith("en") else language
         self.task = task
         self.initial_prompt = initial_prompt
-        self.vad_parameters = vad_parameters or {"threshold": 0.5}
+        self.vad_parameters = vad_parameters or {"onset": 0.5}
         self.no_speech_thresh = 0.45
+        self.same_output_threshold = 10
+        self.end_time_for_same_output = None
 
         device = "cuda" if torch.cuda.is_available() else "cpu"
+        if device == "cuda":
+            major, _ = torch.cuda.get_device_capability(device)
+            self.compute_type = "float16" if major >= 7 else "float32"
+        else:
+            self.compute_type = "int8"
 
         if self.model_size_or_path is None:
             return
+        logging.info(f"Using Device={device} with precision {self.compute_type}")
     
+        try:
             if single_model:
                 if ServeClientFasterWhisper.SINGLE_MODEL is None:
                     self.create_model(device)
@@ -799,6 +822,15 @@ class ServeClientFasterWhisper(ServeClientBase):
                     self.transcriber = ServeClientFasterWhisper.SINGLE_MODEL
             else:
                 self.create_model(device)
+        except Exception as e:
+            logging.error(f"Failed to load model: {e}")
+            self.websocket.send(json.dumps({
+                "uid": self.client_uid,
+                "status": "ERROR",
+                "message": f"Failed to load model: {str(self.model_size_or_path)}"
+            }))
+            self.websocket.close()
+            return
 
         self.use_vad = use_vad
 
@@ -822,7 +854,7 @@ class ServeClientFasterWhisper(ServeClientBase):
         self.transcriber = WhisperModel(
             self.model_size_or_path,
             device=device,
-            compute_type="int8" if device == "cpu" else "float16",
+            compute_type=self.compute_type,
             local_files_only=False,
         )
 
@@ -973,6 +1005,7 @@ class ServeClientFasterWhisper(ServeClientBase):
 
             input_bytes, duration = self.get_audio_chunk_for_processing()
             if duration < 1.0:
+                time.sleep(0.1)     # wait for audio chunks to arrive
                 continue
             try:
                 input_sample = input_bytes.copy()
@@ -988,7 +1021,7 @@ class ServeClientFasterWhisper(ServeClientBase):
                 logging.error(f"[ERROR]: Failed to transcribe audio chunk: {e}")
                 time.sleep(0.01)
 
-    def format_segment(self, start, end, text):
+    def format_segment(self, start, end, text, completed=False):
         """
         Formats a transcription segment with precise start and end times alongside the transcribed text.
 
@@ -1005,7 +1038,8 @@ class ServeClientFasterWhisper(ServeClientBase):
         return {
             'start': "{:.3f}".format(start),
             'end': "{:.3f}".format(end),
-            'text': text
+            'text': text,
+            'completed': completed
         }
 
     def update_segments(self, segments, duration):
@@ -1031,11 +1065,14 @@ class ServeClientFasterWhisper(ServeClientBase):
         """
         offset = None
         self.current_out = ''
+        last_segment = None
+
         # process complete segments
-        if len(segments) > 1:
+        if len(segments) > 1 and segments[-1].no_speech_prob <= self.no_speech_thresh:
             for i, s in enumerate(segments[:-1]):
                 text_ = s.text
                 self.text.append(text_)
+                with self.lock:
                     start, end = self.timestamp_offset + s.start, self.timestamp_offset + min(duration, s.end)
 
                 if start >= end:
@@ -1043,40 +1080,55 @@ class ServeClientFasterWhisper(ServeClientBase):
                 if s.no_speech_prob > self.no_speech_thresh:
                     continue
 
-                self.transcript.append(self.format_segment(start, end, text_))
+                self.transcript.append(self.format_segment(start, end, text_, completed=True))
                 offset = min(duration, s.end)
 
+        # only process the last segment if it satisfies the no_speech_thresh
+        if segments[-1].no_speech_prob <= self.no_speech_thresh:
             self.current_out += segments[-1].text
+            with self.lock:
                 last_segment = self.format_segment(
                     self.timestamp_offset + segments[-1].start,
                     self.timestamp_offset + min(duration, segments[-1].end),
-            self.current_out
+                    self.current_out,
+                    completed=False
                 )
 
+        if self.current_out.strip() == self.prev_out.strip() and self.current_out != '':
+            self.same_output_count += 1
+
+            # if we remove the audio because of same output on the nth reptition we might remove the 
+            # audio thats not yet transcribed so, capturing the time when it was repeated for the first time
+            if self.end_time_for_same_output is None:
+                self.end_time_for_same_output = segments[-1].end
+            time.sleep(0.1)     # wait for some voice activity just in case there is an unitended pause from the speaker for better punctuations.
+        else:
+            self.same_output_count = 0
+            self.end_time_for_same_output = None
+
         # if same incomplete segment is seen multiple times then update the offset
         # and append the segment to the list
-        if self.current_out.strip() == self.prev_out.strip() and self.current_out != '':
+        if self.same_output_count > self.same_output_threshold:
-            self.same_output_threshold += 1
-        else:
-            self.same_output_threshold = 0
-
-        if self.same_output_threshold > 5:
             if not len(self.text) or self.text[-1].strip().lower() != self.current_out.strip().lower():
                 self.text.append(self.current_out)
+                with self.lock:
                     self.transcript.append(self.format_segment(
                         self.timestamp_offset,
-                    self.timestamp_offset + duration,
+                        self.timestamp_offset + min(duration, self.end_time_for_same_output),
-                    self.current_out
+                        self.current_out,
+                        completed=True
                     ))
             self.current_out = ''
-            offset = duration
+            offset = min(duration, self.end_time_for_same_output)
-            self.same_output_threshold = 0
+            self.same_output_count = 0
             last_segment = None
+            self.end_time_for_same_output = None
         else:
             self.prev_out = self.current_out
 
         # update offset
         if offset is not None:
+            with self.lock:
                 self.timestamp_offset += offset
 
         return last_segment

whisper_live/transcriber_tensorrt.py

@@ -1,5 +1,6 @@
 import json
 import re
+import math
 from collections import OrderedDict
 from pathlib import Path
 from typing import Union
@@ -14,7 +15,8 @@ import tensorrt_llm
 import tensorrt_llm.logger as logger
 from tensorrt_llm._utils import (str_dtype_to_torch, str_dtype_to_trt,
                                  trt_dtype_to_torch)
-from tensorrt_llm.runtime import ModelConfig, SamplingConfig
+from tensorrt_llm.bindings import GptJsonConfig, KVCacheType
+from tensorrt_llm.runtime import PYTHON_BINDINGS, ModelConfig, SamplingConfig
 from tensorrt_llm.runtime.session import Session, TensorInfo
 
 
@@ -24,49 +26,101 @@ HOP_LENGTH = 160
 CHUNK_LENGTH = 30
 N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE  # 480000 samples in a 30-second chunk
 
+def read_config(component, engine_dir):
+    config_path = engine_dir / component / 'config.json'
+    with open(config_path, 'r') as f:
+        config = json.load(f)
+    model_config = OrderedDict()
+    model_config.update(config['pretrained_config'])
+    model_config.update(config['build_config'])
+    return model_config
+
+
+def remove_tensor_padding(input_tensor,
+                          input_tensor_lengths=None,
+                          pad_value=None):
+    if pad_value:
+        assert input_tensor_lengths is None, "input_tensor_lengths should be None when pad_value is provided"
+        # Text tensor case: batch, seq_len
+        assert torch.all(
+            input_tensor[:, 0] != pad_value
+        ), "First token in each sequence should not be pad_value"
+        assert input_tensor_lengths is None
+
+        # Create a mask for all non-pad tokens
+        mask = input_tensor != pad_value
+
+        # Apply the mask to input_tensor to remove pad tokens
+        output_tensor = input_tensor[mask].view(1, -1)
+
+    else:
+        # Audio tensor case: batch, seq_len, feature_len
+        # position_ids case: batch, seq_len
+        assert input_tensor_lengths is not None, "input_tensor_lengths must be provided for 3D input_tensor"
+
+        # Initialize a list to collect valid sequences
+        valid_sequences = []
+
+        for i in range(input_tensor.shape[0]):
+            valid_length = input_tensor_lengths[i]
+            valid_sequences.append(input_tensor[i, :valid_length])
+
+        # Concatenate all valid sequences along the batch dimension
+        output_tensor = torch.cat(valid_sequences, dim=0)
+    return output_tensor
+
 
 class WhisperEncoding:
 
     def __init__(self, engine_dir):
         self.session = self.get_session(engine_dir)
+        config = read_config('encoder', engine_dir)
+        self.n_mels = config['n_mels']
+        self.dtype = config['dtype']
+        self.num_languages = config['num_languages']
+        self.encoder_config = config
 
     def get_session(self, engine_dir):
-        config_path = engine_dir / 'encoder_config.json'
+        serialize_path = engine_dir / 'encoder' / 'rank0.engine'
-        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']
-
-        self.dtype = dtype
-        self.n_mels = n_mels
-        self.num_languages = num_languages
-
-        serialize_path = engine_dir / f'whisper_encoder_{self.dtype}_tp1_rank0.engine'
-
         with open(serialize_path, 'rb') as f:
             session = Session.from_serialized_engine(f.read())
-
         return session
 
-    def get_audio_features(self, mel):
+    def get_audio_features(self,
-
+                           mel,
-        input_lengths = torch.tensor(
+                           mel_input_lengths,
-            [mel.shape[2] // 2 for _ in range(mel.shape[0])],
+                           encoder_downsampling_factor=2):
+        if isinstance(mel, list):
+            longest_mel = max([f.shape[-1] for f in mel])
+            mel = [
+                torch.nn.functional.pad(f, (0, longest_mel - f.shape[-1]),
+                                        mode='constant') for f in mel
+            ]
+            mel = torch.cat(mel, dim=0).type(
+                str_dtype_to_torch("float16")).contiguous()
+        bsz, seq_len = mel.shape[0], mel.shape[2]
+        position_ids = torch.arange(
+            math.ceil(seq_len / encoder_downsampling_factor),
             dtype=torch.int32,
-            device=mel.device)
+            device=mel.device).expand(bsz, -1).contiguous()
-
+        if self.encoder_config['plugin_config']['remove_input_padding']:
+            # mel B,D,T -> B,T,D -> BxT, D
+            mel = mel.transpose(1, 2)
+            mel = remove_tensor_padding(mel, mel_input_lengths)
+            position_ids = remove_tensor_padding(
+                position_ids, mel_input_lengths // encoder_downsampling_factor)
         inputs = OrderedDict()
-        inputs['x'] = mel
+        inputs['input_features'] = mel
-        inputs['input_lengths'] = input_lengths
+        inputs['input_lengths'] = mel_input_lengths
+        inputs['position_ids'] = position_ids
 
         output_list = [
-            TensorInfo('x', str_dtype_to_trt(self.dtype), mel.shape),
+            TensorInfo('input_features', str_dtype_to_trt(self.dtype),
+                       mel.shape),
             TensorInfo('input_lengths', str_dtype_to_trt('int32'),
-                       input_lengths.shape)
+                       mel_input_lengths.shape),
+            TensorInfo('position_ids', str_dtype_to_trt('int32'),
+                       inputs['position_ids'].shape)
         ]
 
         output_info = (self.session).infer_shapes(output_list)
@@ -84,48 +138,44 @@ class WhisperEncoding:
                               stream=stream.cuda_stream)
         assert ok, 'Engine execution failed'
         stream.synchronize()
-        audio_features = outputs['output']
+        encoder_output = outputs['encoder_output']
-        return audio_features
+        encoder_output_lengths = mel_input_lengths // encoder_downsampling_factor
+        return encoder_output, encoder_output_lengths
 
 
 class WhisperDecoding:
 
     def __init__(self, engine_dir, runtime_mapping, debug_mode=False):
 
-        self.decoder_config = self.get_config(engine_dir)
+        self.decoder_config = read_config('decoder', engine_dir)
         self.decoder_generation_session = self.get_session(
             engine_dir, runtime_mapping, debug_mode)
 
-    def get_config(self, engine_dir):
-        config_path = engine_dir / 'decoder_config.json'
-        with open(config_path, 'r') as f:
-            config = json.load(f)
-        decoder_config = OrderedDict()
-        decoder_config.update(config['plugin_config'])
-        decoder_config.update(config['builder_config'])
-        return decoder_config
-
     def get_session(self, engine_dir, runtime_mapping, debug_mode=False):
-        dtype = self.decoder_config['precision']
+        serialize_path = engine_dir / 'decoder' / 'rank0.engine'
-        serialize_path = engine_dir / f'whisper_decoder_{dtype}_tp1_rank0.engine'
         with open(serialize_path, "rb") as f:
             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_heads=self.decoder_config['num_attention_heads'],
-            num_kv_heads=self.decoder_config['num_heads'],
+            num_kv_heads=self.decoder_config['num_attention_heads'],
             hidden_size=self.decoder_config['hidden_size'],
             vocab_size=self.decoder_config['vocab_size'],
-            num_layers=self.decoder_config['num_layers'],
+            cross_attention=True,
-            gpt_attention_plugin=self.decoder_config['gpt_attention_plugin'],
+            num_layers=self.decoder_config['num_hidden_layers'],
-            remove_input_padding=self.decoder_config['remove_input_padding'],
+            gpt_attention_plugin=self.decoder_config['plugin_config']
-            cross_attention=self.decoder_config['cross_attention'],
+            ['gpt_attention_plugin'],
+            remove_input_padding=self.decoder_config['plugin_config']
+            ['remove_input_padding'],
+            kv_cache_type=KVCacheType.PAGED
+            if self.decoder_config['plugin_config']['paged_kv_cache'] == True
+            else KVCacheType.CONTINUOUS,
             has_position_embedding=self.
             decoder_config['has_position_embedding'],
-            has_token_type_embedding=self.
+            dtype=self.decoder_config['dtype'],
-            decoder_config['has_token_type_embedding'],
+            has_token_type_embedding=False,
         )
         decoder_generation_session = tensorrt_llm.runtime.GenerationSession(
             decoder_model_config,
@@ -138,14 +188,12 @@ class WhisperDecoding:
     def generate(self,
                  decoder_input_ids,
                  encoder_outputs,
+                 encoder_max_input_length,
+                 encoder_input_lengths,
                  eot_id,
                  max_new_tokens=40,
                  num_beams=1):
-        encoder_input_lengths = torch.tensor(
+        batch_size = decoder_input_ids.shape[0]
-            [encoder_outputs.shape[1] for x in range(encoder_outputs.shape[0])],
-            dtype=torch.int32,
-            device='cuda')
-
         decoder_input_lengths = torch.tensor([
             decoder_input_ids.shape[-1]
             for _ in range(decoder_input_ids.shape[0])
@@ -154,10 +202,10 @@ class WhisperDecoding:
                                              device='cuda')
         decoder_max_input_length = torch.max(decoder_input_lengths).item()
 
-        cross_attention_mask = torch.ones(
+        cross_attention_mask = torch.ones([
-            [encoder_outputs.shape[0], 1,
+            batch_size, decoder_max_input_length + max_new_tokens,
-             encoder_outputs.shape[1]]).int().cuda()
+            encoder_max_input_length
-
+        ]).int().cuda()
         # generation config
         sampling_config = SamplingConfig(end_id=eot_id,
                                          pad_id=eot_id,
@@ -167,11 +215,24 @@ class WhisperDecoding:
             decoder_max_input_length,
             max_new_tokens,
             beam_width=num_beams,
-            encoder_max_input_length=encoder_outputs.shape[1])
+            encoder_max_input_length=encoder_max_input_length)
 
         torch.cuda.synchronize()
 
         decoder_input_ids = decoder_input_ids.type(torch.int32).cuda()
+        if self.decoder_config['plugin_config']['remove_input_padding']:
+            # 50256 is the index of <pad> for all whisper models' decoder
+            WHISPER_PAD_TOKEN_ID = 50256
+            decoder_input_ids = remove_tensor_padding(
+                decoder_input_ids, pad_value=WHISPER_PAD_TOKEN_ID)
+            if encoder_outputs.dim() == 3:
+                encoder_output_lens = torch.full((encoder_outputs.shape[0], ),
+                                                 encoder_outputs.shape[1],
+                                                 dtype=torch.int32,
+                                                 device='cuda')
+
+                encoder_outputs = remove_tensor_padding(encoder_outputs,
+                                                        encoder_output_lens)
         output_ids = self.decoder_generation_session.decode(
             decoder_input_ids,
             decoder_input_lengths,
@@ -196,6 +257,11 @@ class WhisperTRTLLM(object):
         runtime_mapping = tensorrt_llm.Mapping(world_size, runtime_rank)
         torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node)
         engine_dir = Path(engine_dir)
+        encoder_config = read_config('encoder', engine_dir)
+        decoder_config = read_config('decoder', engine_dir)
+        self.n_mels = encoder_config['n_mels']
+        self.num_languages = encoder_config['num_languages']
+        is_multilingual = (decoder_config['vocab_size'] >= 51865)
 
         self.encoder = WhisperEncoding(engine_dir)
         self.decoder = WhisperDecoding(engine_dir,
@@ -207,7 +273,7 @@ class WhisperTRTLLM(object):
         self.device = device
         self.tokenizer = get_tokenizer(
             is_multilingual,
-            num_languages=self.encoder.num_languages,
+            num_languages=self.num_languages,
             language=language,
             task=task,
         )
@@ -274,8 +340,10 @@ class WhisperTRTLLM(object):
     def process_batch(
             self,
             mel,
+            mel_input_lengths,
             text_prefix="<|startoftranscript|><|en|><|transcribe|><|notimestamps|>",
-            num_beams=1):
+            num_beams=1,
+            max_new_tokens=96):
         prompt_id = self.tokenizer.encode(
             text_prefix, allowed_special=set(self.tokenizer.special_tokens.keys()))
 
@@ -283,11 +351,14 @@ class WhisperTRTLLM(object):
         batch_size = mel.shape[0]
         decoder_input_ids = prompt_id.repeat(batch_size, 1)
 
-        encoder_output = self.encoder.get_audio_features(mel)
+        encoder_output, encoder_output_lengths = self.encoder.get_audio_features(mel, mel_input_lengths)
+        encoder_max_input_length = torch.max(encoder_output_lengths).item()
         output_ids = self.decoder.generate(decoder_input_ids,
                                            encoder_output,
+                                           encoder_max_input_length,
+                                           encoder_output_lengths,
                                            self.tokenizer.eot,
-                                           max_new_tokens=96,
+                                           max_new_tokens=max_new_tokens,
                                            num_beams=num_beams)
         texts = []
         for i in range(len(output_ids)):
@@ -302,10 +373,22 @@ class WhisperTRTLLM(object):
             dtype='float16',
             batch_size=1,
             num_beams=1,
+            padding_strategy="max",
             ):
         mel = mel.type(str_dtype_to_torch(dtype))
         mel = mel.unsqueeze(0)
-        predictions = self.process_batch(mel, text_prefix, num_beams)
+        # repeat the mel spectrogram to match the batch size
+        mel = mel.repeat(batch_size, 1, 1)
+        if padding_strategy == "longest":
+            pass
+        else:
+            mel = torch.nn.functional.pad(mel, (0, 3000 - mel.shape[2]))
+        features_input_lengths = torch.full((mel.shape[0], ),
+                                             mel.shape[2],
+                                             dtype=torch.int32,
+                                             device=mel.device)
+
+        predictions = self.process_batch(mel, features_input_lengths, text_prefix, num_beams)
         prediction = predictions[0]
 
         # remove all special tokens in the prediction

whisper_live/vad.py

@@ -1,10 +1,9 @@
-# original: https://github.com/snakers4/silero-vad/blob/master/utils_vad.py
-
 import os
 import subprocess
 import torch
 import numpy as np
 import onnxruntime
+import warnings
 
 
 class VoiceActivityDetection():
@@ -24,6 +23,10 @@ class VoiceActivityDetection():
             self.session = onnxruntime.InferenceSession(path, providers=['CUDAExecutionProvider'], sess_options=opts)
 
         self.reset_states()
+        if '16k' in path:
+            warnings.warn('This model support only 16000 sampling rate!')
+            self.sample_rates = [16000]
+        else:
             self.sample_rates = [8000, 16000]
 
     def _validate_input(self, x, sr: int):
@@ -39,22 +42,27 @@ class VoiceActivityDetection():
 
         if sr not in self.sample_rates:
             raise ValueError(f"Supported sampling rates: {self.sample_rates} (or multiply of 16000)")
-
         if sr / x.shape[1] > 31.25:
             raise ValueError("Input audio chunk is too short")
 
         return x, sr
 
     def reset_states(self, batch_size=1):
-        self._h = np.zeros((2, batch_size, 64)).astype('float32')
+        self._state = torch.zeros((2, batch_size, 128)).float()
-        self._c = np.zeros((2, batch_size, 64)).astype('float32')
+        self._context = torch.zeros(0)
         self._last_sr = 0
         self._last_batch_size = 0
 
     def __call__(self, x, sr: int):
 
         x, sr = self._validate_input(x, sr)
+        num_samples = 512 if sr == 16000 else 256
+
+        if x.shape[-1] != num_samples:
+            raise ValueError(f"Provided number of samples is {x.shape[-1]} (Supported values: 256 for 8000 sample rate, 512 for 16000)")
+
         batch_size = x.shape[0]
+        context_size = 64 if sr == 16000 else 32
 
         if not self._last_batch_size:
             self.reset_states(batch_size)
@@ -63,28 +71,35 @@ class VoiceActivityDetection():
         if (self._last_batch_size) and (self._last_batch_size != batch_size):
             self.reset_states(batch_size)
 
+        if not len(self._context):
+            self._context = torch.zeros(batch_size, context_size)
+
+        x = torch.cat([self._context, x], dim=1)
         if sr in [8000, 16000]:
-            ort_inputs = {'input': x.numpy(), 'h': self._h, 'c': self._c, 'sr': np.array(sr, dtype='int64')}
+            ort_inputs = {'input': x.numpy(), 'state': self._state.numpy(), 'sr': np.array(sr, dtype='int64')}
             ort_outs = self.session.run(None, ort_inputs)
-            out, self._h, self._c = ort_outs
+            out, state = ort_outs
+            self._state = torch.from_numpy(state)
         else:
             raise ValueError()
 
+        self._context = x[..., -context_size:]
         self._last_sr = sr
         self._last_batch_size = batch_size
 
-        out = torch.tensor(out)
+        out = torch.from_numpy(out)
         return out
 
-    def audio_forward(self, x, sr: int, num_samples: int = 512):
+    def audio_forward(self, x, sr: int):
         outs = []
         x, sr = self._validate_input(x, sr)
+        self.reset_states()
+        num_samples = 512 if sr == 16000 else 256
 
         if x.shape[1] % num_samples:
             pad_num = num_samples - (x.shape[1] % num_samples)
             x = torch.nn.functional.pad(x, (0, pad_num), 'constant', value=0.0)
 
-        self.reset_states(x.shape[0])
         for i in range(0, x.shape[1], num_samples):
             wavs_batch = x[:, i:i+num_samples]
             out_chunk = self.__call__(wavs_batch, sr)
@@ -94,7 +109,7 @@ class VoiceActivityDetection():
         return stacked.cpu()
 
     @staticmethod
-    def download(model_url="https://github.com/snakers4/silero-vad/raw/v4.0/files/silero_vad.onnx"):
+    def download(model_url="https://github.com/snakers4/silero-vad/raw/v5.0/files/silero_vad.onnx"):
         target_dir = os.path.expanduser("~/.cache/whisper-live/")
 
         # Ensure the target directory exists
@@ -138,5 +153,5 @@ class VoiceActivityDetector:
             bool: True if the speech probability exceeds the threshold, indicating the presence of voice activity;
                   False otherwise.
         """
-        speech_prob = self.model(torch.from_numpy(audio_frame), self.frame_rate).item()
+        speech_probs = self.model.audio_forward(torch.from_numpy(audio_frame.copy()), self.frame_rate)[0]
-        return speech_prob > self.threshold
+        return torch.any(speech_probs > self.threshold).item()