add eos to faster_whisper server

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

.github/workflows/ci.yml

@@ -15,7 +15,7 @@ jobs:
     runs-on: ubuntu-22.04
     strategy:
       matrix:
-        python-version: [3.8, 3.9, '3.10', 3.11, 3.12]
+        python-version: [3.8, 3.9, '3.10', 3.11]
     steps:
       - uses: actions/checkout@v2
       
@@ -35,7 +35,7 @@ jobs:
             ${{ runner.os }}-pip-${{ matrix.python-version }}-
       
       - name: Install system dependencies
-        run: sudo apt-get update && sudo apt-get install -y portaudio19-dev
+        run: sudo apt-get update && sudo apt-get install -y ffmpeg portaudio19-dev
       
       - name: Install Python dependencies
         run: |
@@ -52,7 +52,7 @@ jobs:
     runs-on: ubuntu-22.04
     strategy:
       matrix:
-        python-version: [3.8, 3.9, '3.10', 3.11, 3.12]
+        python-version: [3.8, 3.9, '3.10', 3.11]
 
     steps:
       - uses: actions/checkout@v2
@@ -101,7 +101,7 @@ jobs:
   
   build-and-push-docker-tensorrt:
     needs: [run-tests, check-code-format]
-    timeout-minutes: 60
+    timeout-minutes: 20
     runs-on: ubuntu-22.04
     if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/'))
     steps:
@@ -180,7 +180,7 @@ jobs:
             ubuntu-latest-pip-3.8-
 
       - name: Install system dependencies
-        run: sudo apt-get update && sudo apt-get install -y portaudio19-dev
+        run: sudo apt-get update && sudo apt-get install -y ffmpeg portaudio19-dev
 
       - name: Install Python dependencies
         run: |

README.md

@@ -12,7 +12,7 @@ to convert speech input into text output. It can be used to transcribe both live
 input from microphone and pre-recorded audio files.
 
 ## Installation
-- Install PyAudio
+- Install PyAudio and ffmpeg
 ```bash
  bash scripts/setup.sh
 ```
@@ -77,10 +77,6 @@ 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.
-  - `mute_audio_playback`: Whether to mute audio playback when transcribing an audio file. Defaults to False.
-
 ```python
 from whisper_live.client import TranscriptionClient
 client = TranscriptionClient(
@@ -88,13 +84,10 @@ client = TranscriptionClient(
   9090,
   lang="en",
   translate=False,
-  model="small",                                      # also support hf_model => `Systran/faster-whisper-small`
+  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
-  max_clients=4,
-  max_connection_time=600,
-  mute_audio_playback=False,                          # Only used for file 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.
@@ -134,17 +127,13 @@ 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 small.en engine
-  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
+  # Build tiny.en engine
+  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en
 
-  # Run server with 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_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"
+                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en"
   ```
 
 - CPU

TensorRT_whisper.md

@@ -1,11 +1,17 @@
 # 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.15.0.dev2024111200`
+**Note**: We use `tensorrt_llm==0.9.0`
 
 ## 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
@@ -15,9 +21,7 @@ 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        # 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
+bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en
 
 # convert small multilingual model
 bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small
@@ -28,11 +32,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_float16"
+                      --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 "/app/TensorRT-LLM-examples/whisper/whisper_small_float16" \
+                      --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small" \
                       --trt_multilingual
 ```

docker/Dockerfile.tensorrt

@@ -1,22 +1,19 @@
-FROM nvidia/cuda:12.4.1-base-ubuntu22.04 AS base
+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 git-lfs wget \
+    python3.10 python3-pip openmpi-bin libopenmpi-dev git wget \
     && rm -rf /var/lib/apt/lists/*
 
-FROM base AS devel
-RUN pip3 install --no-cache-dir -U tensorrt_llm==0.15.0.dev2024111200 --extra-index-url https://pypi.nvidia.com
+RUN pip3 install --no-cache-dir -U tensorrt_llm==0.9.0 --extra-index-url https://pypi.nvidia.com
+
 WORKDIR /app
-RUN git clone https://github.com/NVIDIA/TensorRT-LLM.git && cd TensorRT-LLM && \
-    git checkout c629546ce429623c8a163633095230154a6f0574 && cd ../ && \
+
+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
 
-
-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
 
@@ -25,7 +22,7 @@ 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
-RUN pip install pynvml==11.5.0
+
 COPY whisper_live ./whisper_live
 COPY scripts/build_whisper_tensorrt.sh .
 COPY run_server.py .

requirements/client.txt

@@ -1,4 +1,4 @@
 PyAudio
-av
+ffmpeg-python
 scipy
 websocket-client

requirements/server.txt

@@ -1,13 +1,13 @@
-faster-whisper==1.1.0
+faster-whisper==1.0.1
+torch
 websockets
-onnxruntime==1.17.0
+onnxruntime==1.16.0
 numba
+openai-whisper
 kaldialign
 soundfile
+ffmpeg-python
 scipy
-av
 jiwer
 evaluate
 numpy<2
-openai-whisper==20240930
-tokenizers==0.20.3

scripts/build_whisper_tensorrt.sh

@@ -38,24 +38,12 @@ 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"
@@ -66,43 +54,11 @@ download_and_build_model() {
         echo "${model_name}.pt already exists in assets directory."
     fi
 
-    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}"
+    local output_dir="whisper_${model_name//./_}"
     echo "$output_dir"
-    echo "Converting model weights for $model_name..."
-    python3 convert_checkpoint.py \
-        $( [[ "$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 "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 --enable_context_fmha --model_name "$model_name"
+    echo "Whisper $model_name TensorRT engine built."
     echo "========================================="
     echo "Model is located at: $(pwd)/$output_dir"
 }
@@ -114,9 +70,8 @@ fi
 
 tensorrt_examples_dir="$1"
 model_name="${2:-small.en}"
-weight_only_precision="${3:-float16}"  # Default to float16 if not provided
 
-cd $tensorrt_examples_dir/whisper
+cd $1/whisper
 pip install --no-deps -r requirements.txt
 
-download_and_build_model "$model_name" "$weight_only_precision"
+download_and_build_model "$model_name"

scripts/setup.sh

@@ -1,3 +1,3 @@
 #! /bin/bash
 
-apt-get install portaudio19-dev wget -y
+apt-get install portaudio19-dev ffmpeg wget -y

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,18 +43,18 @@ setup(
     ),
     install_requires=[
         "PyAudio",
-        "faster-whisper==1.1.0",
+        "faster-whisper==1.0.1",
         "torch",
         "torchaudio",
         "websockets",
-        "onnxruntime==1.17.0",
+        "onnxruntime==1.16.0",
+        "ffmpeg-python",
         "scipy",
         "websocket-client",
         "numba",
-        "openai-whisper==20240930",
+        "openai-whisper",
         "kaldialign",
         "soundfile",
-        "tokenizers==0.20.3"
     ],
     python_requires=">=3.8"
 )

tests/test_client.py

@@ -48,9 +48,7 @@ class TestClientCallbacks(BaseTestCase):
             "language": self.client.language,
             "task": self.client.task,
             "model": self.client.model,
-            "use_vad": True,
-            "max_clients": 4,
-            "max_connection_time": 600,
+            "use_vad": True
         })
         self.client.on_open(self.mock_ws_app)
         self.mock_ws_app.send.assert_called_with(expected_message)
@@ -68,15 +66,15 @@ class TestClientCallbacks(BaseTestCase):
         message = json.dumps({
             "uid": self.client.uid,
             "segments": [
-                {"start": 0, "end": 1, "text": "Test transcript", "completed": True},
-                {"start": 1, "end": 2, "text": "Test transcript 2", "completed": True},
-                {"start": 2, "end": 3, "text": "Test transcript 3", "completed": True}
+                {"start": 0, "end": 1, "text": "Test transcript"},
+                {"start": 1, "end": 2, "text": "Test transcript 2"},
+                {"start": 2, "end": 3, "text": "Test transcript 3"}
             ]
         })
         self.client.on_message(self.mock_ws_app, message)
 
         # Assert that the transcript was updated correctly
-        self.assertEqual(len(self.client.transcript), 3)
+        self.assertEqual(len(self.client.transcript), 2)
         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 jiwer
+import evaluate
 
 from websockets.exceptions import ConnectionClosed
-from whisper_live.server import TranscriptionServer, BackendType, ClientManager
+from whisper_live.server import TranscriptionServer
 from whisper_live.client import Client, TranscriptionClient, TranscriptionTeeClient
 from whisper.normalizers import EnglishTextNormalizer
 
@@ -16,7 +16,6 @@ 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, {})
@@ -26,7 +25,6 @@ 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
@@ -51,7 +49,7 @@ class TestServerConnection(unittest.TestCase):
             'task': 'transcribe',
             'model': 'tiny.en'
         })
-        self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
+        self.server.recv_audio(mock_websocket, "faster_whisper")
 
     @mock.patch('websockets.WebSocketCommonProtocol')
     def test_recv_audio_exception_handling(self, mock_websocket):
@@ -63,7 +61,7 @@ class TestServerConnection(unittest.TestCase):
         }),  np.array([1, 2, 3]).tobytes()]
 
         with self.assertLogs(level="ERROR"):
-            self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
+            self.server.recv_audio(mock_websocket, "faster_whisper")
 
         self.assertNotIn(mock_websocket, self.server.client_manager.clients)
 
@@ -84,6 +82,7 @@ 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):
@@ -95,8 +94,11 @@ class TestServerInferenceAccuracy(unittest.TestCase):
         gt_normalized = self.normalizer(gt)
 
         # calculate WER
-        wer_score = jiwer.wer(gt_normalized, prediction_normalized)
-        self.assertLess(wer_score, 0.05)
+        wer = self.metric.compute(
+            predictions=[prediction_normalized],
+            references=[gt_normalized]
+        )
+        self.assertLess(wer, 0.05)
 
     def test_inference(self):
         client = TranscriptionClient(
@@ -122,10 +124,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", rcvd_then_sent=mock.Mock())
+        mock_websocket.recv.side_effect = ConnectionClosed(1001, "testing connection closed")
 
         with self.assertLogs(level="INFO") as log:
-            self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
+            self.server.recv_audio(mock_websocket, "faster_whisper")
             self.assertTrue(any("Connection closed by client" in message for message in log.output))
 
     @mock.patch('websockets.WebSocketCommonProtocol')
@@ -133,7 +135,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, BackendType("faster_whisper"))
+            self.server.recv_audio(mock_websocket, "faster_whisper")
             self.assertTrue(any("Failed to decode JSON from client" in message for message in log.output))
 
     @mock.patch('websockets.WebSocketCommonProtocol')
@@ -141,7 +143,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, BackendType("faster_whisper"))
+            self.server.recv_audio(mock_websocket, "faster_whisper")
             for message in log.output:
                 print(message)
             print()

whisper_live/__version__.py

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

whisper_live/client.py

@@ -2,7 +2,6 @@ import os
 import shutil
 import wave
 
-import logging
 import numpy as np
 import pyaudio
 import threading
@@ -10,7 +9,7 @@ import json
 import websocket
 import uuid
 import time
-import av
+import ffmpeg
 import whisper_live.utils as utils
 
 
@@ -29,10 +28,7 @@ class Client:
         translate=False,
         model="small",
         srt_file_path="output.srt",
-        use_vad=True,
-        log_transcription=True,
-        max_clients=4,
-        max_connection_time=600,
+        use_vad=True
     ):
         """
         Initializes a Client instance for audio recording and streaming to a server.
@@ -46,12 +42,6 @@ class Client:
             port (int): The port number for the WebSocket server.
             lang (str, optional): The selected language for transcription. Default is None.
             translate (bool, optional): Specifies if the task is translation. Default is False.
-            model (str, optional): The whisper model to use (e.g., "small", "medium", "large"). Default is "small".
-            srt_file_path (str, optional): The file path to save the output SRT file. Default is "output.srt".
-            use_vad (bool, optional): Whether to enable voice activity detection. Default is True.
-            log_transcription (bool, optional): Whether to log transcription output to the console. Default is True.
-            max_clients (int, optional): Maximum number of client connections allowed. Default is 4.
-            max_connection_time (int, optional): Maximum allowed connection time in seconds. Default is 600.
         """
         self.recording = False
         self.task = "transcribe"
@@ -66,13 +56,11 @@ 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:
@@ -118,9 +106,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 and not seg.get("completed", False):
+                if i == len(segments) - 1:
                     self.last_segment = seg
-                elif (self.server_backend == "faster_whisper" and seg.get("completed", False) and
+                elif (self.server_backend == "faster_whisper" and
                       (not self.transcript or
                         float(seg['start']) >= float(self.transcript[-1]['end']))):
                     self.transcript.append(seg)
@@ -129,11 +117,10 @@ 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()
-            utils.print_transcript(text)
+        # Truncate to last 3 entries for brevity.
+        text = text[-3:]
+        utils.clear_screen()
+        utils.print_transcript(text)
 
     def on_message(self, ws, message):
         """
@@ -209,9 +196,7 @@ class Client:
                     "language": self.language,
                     "task": self.task,
                     "model": self.model,
-                    "use_vad": self.use_vad,
-                    "max_clients": self.max_clients,
-                    "max_connection_time": self.max_connection_time,
+                    "use_vad": self.use_vad
                 }
             )
         )
@@ -265,9 +250,7 @@ class Client:
 
         """
         if self.server_backend == "faster_whisper":
-            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"]:
+            if (self.last_segment):
                 self.transcript.append(self.last_segment)
             utils.create_srt_file(self.transcript, output_path)
 
@@ -291,7 +274,7 @@ class TranscriptionTeeClient:
     Attributes:
         clients (list): the underlying Client instances responsible for handling WebSocket connections.
     """
-    def __init__(self, clients, save_output_recording=False, output_recording_filename="./output_recording.wav", mute_audio_playback=False):
+    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.")
@@ -302,7 +285,6 @@ class TranscriptionTeeClient:
         self.record_seconds = 60000
         self.save_output_recording = save_output_recording
         self.output_recording_filename = output_recording_filename
-        self.mute_audio_playback = mute_audio_playback
         self.frames = b""
         self.p = pyaudio.PyAudio()
         try:
@@ -398,7 +380,6 @@ class TranscriptionTeeClient:
                 output=True,
                 frames_per_buffer=self.chunk,
             )
-            chunk_duration = self.chunk / float(wavfile.getframerate())
             try:
                 while any(client.recording for client in self.clients):
                     data = wavfile.readframes(self.chunk)
@@ -407,10 +388,7 @@ class TranscriptionTeeClient:
 
                     audio_array = self.bytes_to_float_array(data)
                     self.multicast_packet(audio_array.tobytes())
-                    if self.mute_audio_playback:
-                        time.sleep(chunk_duration)
-                    else:
-                        self.stream.write(data)
+                    self.stream.write(data)
 
                 wavfile.close()
 
@@ -432,83 +410,72 @@ class TranscriptionTeeClient:
 
     def process_rtsp_stream(self, rtsp_url):
         """
-        Connect to an RTSP source, process the audio stream, and send it for transcription.
+        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.
         """
-        print("[INFO]: Connecting to RTSP stream...")
-        try:
-            container = av.open(rtsp_url, format="rtsp", options={"rtsp_transport": "tcp"})
-            self.process_av_stream(container, stream_type="RTSP")
-        except Exception as e:
-            print(f"[ERROR]: Failed to process RTSP stream: {e}")
-        finally:
-            for client in self.clients:
-                client.wait_before_disconnect()
-            self.multicast_packet(Client.END_OF_AUDIO.encode('utf-8'), True)
-            self.close_all_clients()
-            self.write_all_clients_srt()
-        print("[INFO]: RTSP stream processing finished.")
+        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=None):
+    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.
+            save_file （str, optional): Local path to save the network stream.
         """
-        print("[INFO]: Connecting to HLS stream...")
+        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:
-            container = av.open(hls_url, format="hls")
-            self.process_av_stream(container, stream_type="HLS", save_file=save_file)
+            # Process the stream
+            while True:
+                in_bytes = process.stdout.read(self.chunk * 2)  # 2 bytes per sample
+                if not in_bytes:
+                    break
+                audio_array = self.bytes_to_float_array(in_bytes)
+                self.multicast_packet(audio_array.tobytes())
+
         except Exception as e:
-            print(f"[ERROR]: Failed to process HLS stream: {e}")
+            print(f"[ERROR]: Failed to connect to {stream_type} stream: {e}")
         finally:
-            for client in self.clients:
-                client.wait_before_disconnect()
-            self.multicast_packet(Client.END_OF_AUDIO.encode('utf-8'), True)
             self.close_all_clients()
             self.write_all_clients_srt()
-        print("[INFO]: HLS stream processing finished.")
+            if process:
+                process.kill()
 
-    def process_av_stream(self, container, stream_type, save_file=None):
-        """
-        Process an AV container stream and send audio packets to the server.
+        print(f"[INFO]: {stream_type} stream processing finished.")
 
-        Args:
-            container (av.container.InputContainer): The input container to process.
-            stream_type (str): The type of stream being processed ("RTSP" or "HLS").
-            save_file (str, optional): Local path to save the stream. Default is None.
-        """
-        audio_stream = next((s for s in container.streams if s.type == "audio"), None)
-        if not audio_stream:
-            print(f"[ERROR]: No audio stream found in {stream_type} source.")
-            return
+    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)
+        )
 
-        output_container = None
-        if save_file:
-            output_container = av.open(save_file, mode="w")
-            output_audio_stream = output_container.add_stream(codec_name="pcm_s16le", rate=self.rate)
+    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)
+            )
 
-        try:
-            for packet in container.demux(audio_stream):
-                for frame in packet.decode():
-                    audio_data = frame.to_ndarray().tobytes()
-                    self.multicast_packet(audio_data)
-
-                    if save_file:
-                        output_container.mux(frame)
-        except Exception as e:
-            print(f"[ERROR]: Error during {stream_type} stream processing: {e}")
-        finally:
-            # Wait for server to send any leftover transcription.
-            time.sleep(5)
-            self.multicast_packet(Client.END_OF_AUDIO.encode('utf-8'), True)
-            if output_container:
-                output_container.close()
-            container.close()
+        return process
 
     def save_chunk(self, n_audio_file):
         """
@@ -672,16 +639,10 @@ class TranscriptionClient(TranscriptionTeeClient):
         host (str): The hostname or IP address of the server.
         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): If True, the task will be translation instead of transcription. Default is False.
-        model (str, optional): The whisper model to use (e.g., "small", "base"). Default is "small".
-        use_vad (bool, optional): Whether to enable voice activity detection. Default is True.
-        save_output_recording (bool, optional): Whether to save the microphone recording. Default is False.
-        output_recording_filename (str, optional): Path to save the output recording WAV file. Default is "./output_recording.wav".
-        output_transcription_path (str, optional): File path to save the output transcription (SRT file). Default is "./output.srt".
-        log_transcription (bool, optional): Whether to log transcription output to the console. Default is True.
-        max_clients (int, optional): Maximum number of client connections allowed. Default is 4.
-        max_connection_time (int, optional): Maximum allowed connection time in seconds. Default is 600.
-        mute_audio_playback (bool, optional): If True, mutes audio playback during file playback. Default is False.
+        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.
@@ -703,18 +664,9 @@ class TranscriptionClient(TranscriptionTeeClient):
         use_vad=True,
         save_output_recording=False,
         output_recording_filename="./output_recording.wav",
-        output_transcription_path="./output.srt",
-        log_transcription=True,
-        max_clients=4,
-        max_connection_time=600,
-        mute_audio_playback=False,
+        output_transcription_path="./output.srt"
     ):
-        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
-        )
-
+        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"):
@@ -723,6 +675,5 @@ class TranscriptionClient(TranscriptionTeeClient):
             self,
             [self.client],
             save_output_recording=save_output_recording,
-            output_recording_filename=output_recording_filename,
-            mute_audio_playback=mute_audio_playback
+            output_recording_filename=output_recording_filename
         )

whisper_live/server.py

@@ -147,7 +147,7 @@ class TranscriptionServer:
     RATE = 16000
 
     def __init__(self):
-        self.client_manager = None
+        self.client_manager = ClientManager()
         self.no_voice_activity_chunks = 0
         self.use_vad = True
         self.single_model = False
@@ -181,26 +181,22 @@ 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}")
-                    options["model"] = faster_whisper_custom_model_path
-                client = ServeClientFasterWhisper(
-                    websocket,
-                    language=options["language"],
-                    task=options["task"],
-                    client_uid=options["uid"],
-                    model=options["model"],
-                    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.")
-        except Exception as e:
-            return
+        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
+            client = ServeClientFasterWhisper(
+                websocket,
+                language=options["language"],
+                task=options["task"],
+                client_uid=options["uid"],
+                model=options["model"],
+                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.")
@@ -228,19 +224,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()
                 return False  # Indicates that the connection should not continue
 
-            if self.backend.is_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
@@ -258,17 +247,15 @@ 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.is_tensorrt():
-                client.set_eos(True)
+            client.set_eos(True)
             return False
 
-        if self.backend.is_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
@@ -341,13 +328,8 @@ class TranscriptionServer:
             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.")
-        if single_model:
-            if faster_whisper_custom_model_path or whisper_tensorrt_path:
-                logging.info("Custom model option was provided. Switching to single model mode.")
-                self.single_model = True
-                # TODO: load model initially
-            else:
-                logging.info("Single model mode currently only works with custom models.")
+
+        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(
@@ -421,11 +403,12 @@ class ServeClientBase(object):
         self.prev_out = ''
         self.t_start = None
         self.exit = False
-        self.same_output_count = 0
+        self.same_output_threshold = 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 = []
         self.send_last_n_segments = 10
+        self.eos = False
 
         # text formatting
         self.pick_previous_segments = 2
@@ -433,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
 
@@ -479,10 +474,9 @@ 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
+        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
 
     def get_audio_chunk_for_processing(self):
         """
@@ -498,9 +492,8 @@ 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()
+        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
         return input_bytes, duration
 
@@ -555,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:
@@ -660,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.
@@ -717,9 +700,7 @@ 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
+        self.timestamp_offset += duration
 
     def speech_to_text(self):
         """
@@ -788,49 +769,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", "distil-small.en",
-            "distil-medium.en", "distil-large-v2", "distil-large-v3",
-            "large-v3-turbo", "turbo"
+            "medium", "medium.en", "large-v2", "large-v3",
         ]
-
-        self.model_size_or_path = model
+        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 {"onset": 0.5}
-        self.no_speech_thresh = 0.45
-        self.same_output_threshold = 10
-        self.end_time_for_same_output = None
+        self.vad_parameters = vad_parameters or {"threshold": 0.5}
+        self.no_speech_thresh = 0.35
 
         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)
-                    ServeClientFasterWhisper.SINGLE_MODEL = self.transcriber
-                else:
-                    self.transcriber = ServeClientFasterWhisper.SINGLE_MODEL
-            else:
+        if single_model:
+            if ServeClientFasterWhisper.SINGLE_MODEL is None:
                 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
+                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
 
@@ -854,7 +818,7 @@ class ServeClientFasterWhisper(ServeClientBase):
         self.transcriber = WhisperModel(
             self.model_size_or_path,
             device=device,
-            compute_type=self.compute_type,
+            compute_type="int8" if device == "cpu" else "float16",
             local_files_only=False,
         )
 
@@ -920,8 +884,9 @@ class ServeClientFasterWhisper(ServeClientBase):
             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()
 
@@ -964,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):
         """
@@ -1004,8 +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:
-                time.sleep(0.1)     # wait for audio chunks to arrive
+            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()
@@ -1013,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)
 
@@ -1021,7 +989,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, completed=False):
+    def format_segment(self, start, end, text):
         """
         Formats a transcription segment with precise start and end times alongside the transcribed text.
 
@@ -1038,8 +1006,7 @@ class ServeClientFasterWhisper(ServeClientBase):
         return {
             'start': "{:.3f}".format(start),
             'end': "{:.3f}".format(end),
-            'text': text,
-            'completed': completed
+            'text': text
         }
 
     def update_segments(self, segments, duration):
@@ -1063,72 +1030,56 @@ 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 = ''
-        last_segment = None
-
         # process complete segments
-        if len(segments) > 1 and segments[-1].no_speech_prob <= self.no_speech_thresh:
+        if len(segments) > 1:
             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)
+                start, end = self.timestamp_offset + s.start, self.timestamp_offset + min(duration, s.end)
 
                 if start >= end:
                     continue
                 if s.no_speech_prob > self.no_speech_thresh:
                     continue
 
-                self.transcript.append(self.format_segment(start, end, text_, completed=True))
+                self.text.append(text_)
+                self.transcript.append(self.format_segment(start, end, text_))
                 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,
-                    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
+            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
-        if self.same_output_count > self.same_output_threshold:
+        if self.current_out.strip() == self.prev_out.strip() and self.current_out != '':
+            self.same_output_threshold += 1
+        else:
+            self.same_output_threshold = 0
+
+        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)
-                with self.lock:
-                    self.transcript.append(self.format_segment(
-                        self.timestamp_offset,
-                        self.timestamp_offset + min(duration, self.end_time_for_same_output),
-                        self.current_out,
-                        completed=True
-                    ))
+                self.transcript.append(self.format_segment(
+                    self.timestamp_offset,
+                    self.timestamp_offset + duration,
+                    self.current_out
+                ))
             self.current_out = ''
-            offset = min(duration, self.end_time_for_same_output)
-            self.same_output_count = 0
+            offset = duration
+            self.same_output_threshold = 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
+            self.timestamp_offset += offset
 
         return last_segment

whisper_live/tensorrt_utils.py

@@ -23,12 +23,8 @@ from typing import Dict, Iterable, List, Optional, TextIO, Tuple, Union
 import kaldialign
 import numpy as np
 import soundfile
-import av
-import wave
 import torch
 import torch.nn.functional as F
-from whisper_live.utils import resample
-
 
 Pathlike = Union[str, Path]
 
@@ -39,33 +35,38 @@ CHUNK_LENGTH = 30
 N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE  # 480000 samples in a 30-second chunk
 
 
-def load_audio(file: str, sr: int = 16000):
+def load_audio(file: str, sr: int = SAMPLE_RATE):
     """
-    Open an audio file, resample it, and read as a mono waveform.
+    Open an audio file and read as mono waveform, resampling as necessary
 
     Parameters
     ----------
     file: str
-        The audio file to open.
+        The audio file to open
 
     sr: int
-        The sample rate to resample the audio if necessary.
+        The sample rate to resample the audio if necessary
 
     Returns
     -------
     A NumPy array containing the audio waveform, in float32 dtype.
     """
-    resampled_file = resample(file, sr)
 
-    with wave.open(resampled_file, "rb") as wav_file:
-        num_frames = wav_file.getnframes()
-        raw_data = wav_file.readframes(num_frames)
+    # This launches a subprocess to decode audio while down-mixing
+    # and resampling as necessary.  Requires the ffmpeg CLI in PATH.
+    # fmt: off
+    cmd = [
+        "ffmpeg", "-nostdin", "-threads", "0", "-i", file, "-f", "s16le", "-ac",
+        "1", "-acodec", "pcm_s16le", "-ar",
+        str(sr), "-"
+    ]
+    # fmt: on
+    try:
+        out = run(cmd, capture_output=True, check=True).stdout
+    except CalledProcessError as e:
+        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
 
-        audio_data = np.frombuffer(raw_data, dtype=np.int16)
-
-    audio_data = audio_data.astype(np.float32) / 32768.0
-
-    return audio_data
+    return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
 
 
 def load_audio_wav_format(wav_path):

whisper_live/transcriber_tensorrt.py

@@ -1,6 +1,5 @@
 import json
 import re
-import math
 from collections import OrderedDict
 from pathlib import Path
 from typing import Union
@@ -15,8 +14,7 @@ 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.bindings import GptJsonConfig, KVCacheType
-from tensorrt_llm.runtime import PYTHON_BINDINGS, ModelConfig, SamplingConfig
+from tensorrt_llm.runtime import ModelConfig, SamplingConfig
 from tensorrt_llm.runtime.session import Session, TensorInfo
 
 
@@ -26,101 +24,49 @@ 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):
-        serialize_path = engine_dir / 'encoder' / 'rank0.engine'
+        config_path = engine_dir / 'encoder_config.json'
+        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,
-                           mel_input_lengths,
-                           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),
+    def get_audio_features(self, mel):
+
+        input_lengths = torch.tensor(
+            [mel.shape[2] // 2 for _ in range(mel.shape[0])],
             dtype=torch.int32,
-            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)
+            device=mel.device)
+
         inputs = OrderedDict()
-        inputs['input_features'] = mel
-        inputs['input_lengths'] = mel_input_lengths
-        inputs['position_ids'] = position_ids
+        inputs['x'] = mel
+        inputs['input_lengths'] = input_lengths
 
         output_list = [
-            TensorInfo('input_features', str_dtype_to_trt(self.dtype),
-                       mel.shape),
+            TensorInfo('x', str_dtype_to_trt(self.dtype), mel.shape),
             TensorInfo('input_lengths', str_dtype_to_trt('int32'),
-                       mel_input_lengths.shape),
-            TensorInfo('position_ids', str_dtype_to_trt('int32'),
-                       inputs['position_ids'].shape)
+                       input_lengths.shape)
         ]
 
         output_info = (self.session).infer_shapes(output_list)
@@ -138,44 +84,48 @@ class WhisperEncoding:
                               stream=stream.cuda_stream)
         assert ok, 'Engine execution failed'
         stream.synchronize()
-        encoder_output = outputs['encoder_output']
-        encoder_output_lengths = mel_input_lengths // encoder_downsampling_factor
-        return encoder_output, encoder_output_lengths
+        audio_features = outputs['output']
+        return audio_features
 
 
 class WhisperDecoding:
 
     def __init__(self, engine_dir, runtime_mapping, debug_mode=False):
 
-        self.decoder_config = read_config('decoder', engine_dir)
+        self.decoder_config = self.get_config(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):
-        serialize_path = engine_dir / 'decoder' / 'rank0.engine'
+        dtype = self.decoder_config['precision']
+        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_attention_heads'],
-            num_kv_heads=self.decoder_config['num_attention_heads'],
+            num_heads=self.decoder_config['num_heads'],
+            num_kv_heads=self.decoder_config['num_heads'],
             hidden_size=self.decoder_config['hidden_size'],
             vocab_size=self.decoder_config['vocab_size'],
-            cross_attention=True,
-            num_layers=self.decoder_config['num_hidden_layers'],
-            gpt_attention_plugin=self.decoder_config['plugin_config']
-            ['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,
+            num_layers=self.decoder_config['num_layers'],
+            gpt_attention_plugin=self.decoder_config['gpt_attention_plugin'],
+            remove_input_padding=self.decoder_config['remove_input_padding'],
+            cross_attention=self.decoder_config['cross_attention'],
             has_position_embedding=self.
             decoder_config['has_position_embedding'],
-            dtype=self.decoder_config['dtype'],
-            has_token_type_embedding=False,
+            has_token_type_embedding=self.
+            decoder_config['has_token_type_embedding'],
         )
         decoder_generation_session = tensorrt_llm.runtime.GenerationSession(
             decoder_model_config,
@@ -188,12 +138,14 @@ 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):
-        batch_size = decoder_input_ids.shape[0]
+        encoder_input_lengths = torch.tensor(
+            [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])
@@ -202,10 +154,10 @@ class WhisperDecoding:
                                              device='cuda')
         decoder_max_input_length = torch.max(decoder_input_lengths).item()
 
-        cross_attention_mask = torch.ones([
-            batch_size, decoder_max_input_length + max_new_tokens,
-            encoder_max_input_length
-        ]).int().cuda()
+        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,
@@ -215,24 +167,11 @@ class WhisperDecoding:
             decoder_max_input_length,
             max_new_tokens,
             beam_width=num_beams,
-            encoder_max_input_length=encoder_max_input_length)
+            encoder_max_input_length=encoder_outputs.shape[1])
 
         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,
@@ -257,23 +196,18 @@ 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,
-                                        runtime_mapping,
-                                        debug_mode=False)
+                                       runtime_mapping,
+                                       debug_mode=False)
         self.n_mels = self.encoder.n_mels
         # self.tokenizer = get_tokenizer(num_languages=self.encoder.num_languages,
         #                                tokenizer_dir=assets_dir)
         self.device = device
         self.tokenizer = get_tokenizer(
             is_multilingual,
-            num_languages=self.num_languages,
+            num_languages=self.encoder.num_languages,
             language=language,
             task=task,
         )
@@ -340,10 +274,8 @@ class WhisperTRTLLM(object):
     def process_batch(
             self,
             mel,
-            mel_input_lengths,
             text_prefix="<|startoftranscript|><|en|><|transcribe|><|notimestamps|>",
-            num_beams=1,
-            max_new_tokens=96):
+            num_beams=1):
         prompt_id = self.tokenizer.encode(
             text_prefix, allowed_special=set(self.tokenizer.special_tokens.keys()))
 
@@ -351,14 +283,11 @@ class WhisperTRTLLM(object):
         batch_size = mel.shape[0]
         decoder_input_ids = prompt_id.repeat(batch_size, 1)
 
-        encoder_output, encoder_output_lengths = self.encoder.get_audio_features(mel, mel_input_lengths)
-        encoder_max_input_length = torch.max(encoder_output_lengths).item()
+        encoder_output = self.encoder.get_audio_features(mel)
         output_ids = self.decoder.generate(decoder_input_ids,
                                            encoder_output,
-                                           encoder_max_input_length,
-                                           encoder_output_lengths,
                                            self.tokenizer.eot,
-                                           max_new_tokens=max_new_tokens,
+                                           max_new_tokens=96,
                                            num_beams=num_beams)
         texts = []
         for i in range(len(output_ids)):
@@ -373,22 +302,10 @@ 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)
-        # 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)
+        predictions = self.process_batch(mel, text_prefix, num_beams)
         prediction = predictions[0]
 
         # remove all special tokens in the prediction

whisper_live/utils.py

@@ -1,9 +1,8 @@
 import os
 import textwrap
 import scipy
+import ffmpeg
 import numpy as np
-import av
-from pathlib import Path
 
 
 def clear_screen():
@@ -27,8 +26,8 @@ def format_time(s):
     return f"{hours:02}:{minutes:02}:{seconds:02},{milliseconds:03}"
 
 
-def create_srt_file(segments, resampled_file):
-    with open(resampled_file, 'w', encoding='utf-8') as srt_file:
+def create_srt_file(segments, output_file):
+    with open(output_file, 'w', encoding='utf-8') as srt_file:
         segment_number = 1
         for segment in segments:
             start_time = format_time(float(segment['start']))
@@ -44,7 +43,9 @@ def create_srt_file(segments, resampled_file):
 
 def resample(file: str, sr: int = 16000):
     """
-    Resample the audio file to 16kHz.
+    # https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/audio.py#L22
+    Open an audio file and read as mono waveform, resampling as necessary,
+    save the resampled audio
 
     Args:
         file (str): The audio file to open
@@ -53,30 +54,18 @@ def resample(file: str, sr: int = 16000):
     Returns:
         resampled_file (str): The resampled audio file
     """
-    container = av.open(file)
-    stream = next(s for s in container.streams if s.type == 'audio')
+    try:
+        # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
+        # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
+        out, _ = (
+            ffmpeg.input(file, threads=0)
+            .output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=sr)
+            .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
+        )
+    except ffmpeg.Error as e:
+        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
+    np_buffer = np.frombuffer(out, dtype=np.int16)
 
-    resampler = av.AudioResampler(
-        format='s16',
-        layout='mono',
-        rate=sr,
-    )
-
-    resampled_file = Path(file).stem + "_resampled.wav"
-    output_container = av.open(resampled_file, mode='w')
-    output_stream = output_container.add_stream('pcm_s16le', rate=sr)
-    output_stream.layout = 'mono'
-
-    for frame in container.decode(audio=0):
-        frame.pts = None
-        resampled_frames = resampler.resample(frame)
-        if resampled_frames is not None:
-            for resampled_frame in resampled_frames:
-                for packet in output_stream.encode(resampled_frame):
-                    output_container.mux(packet)
-
-    for packet in output_stream.encode(None):
-        output_container.mux(packet)
-
-    output_container.close()
+    resampled_file = f"{file.split('.')[0]}_resampled.wav"
+    scipy.io.wavfile.write(resampled_file, sr, np_buffer.astype(np.int16))
     return resampled_file

whisper_live/vad.py

@@ -1,9 +1,10 @@
+# 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():
@@ -23,11 +24,7 @@ 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]
+        self.sample_rates = [8000, 16000]
 
     def _validate_input(self, x, sr: int):
         if x.dim() == 1:
@@ -37,32 +34,27 @@ class VoiceActivityDetection():
 
         if sr != 16000 and (sr % 16000 == 0):
             step = sr // 16000
-            x = x[:,::step]
+            x = x[:, ::step]
             sr = 16000
 
         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._state = torch.zeros((2, batch_size, 128)).float()
-        self._context = torch.zeros(0)
+        self._h = np.zeros((2, batch_size, 64)).astype('float32')
+        self._c = np.zeros((2, batch_size, 64)).astype('float32')
         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)
@@ -71,35 +63,28 @@ 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(), 'state': self._state.numpy(), 'sr': np.array(sr, dtype='int64')}
+            ort_inputs = {'input': x.numpy(), 'h': self._h, 'c': self._c, 'sr': np.array(sr, dtype='int64')}
             ort_outs = self.session.run(None, ort_inputs)
-            out, state = ort_outs
-            self._state = torch.from_numpy(state)
+            out, self._h, self._c = ort_outs
         else:
             raise ValueError()
 
-        self._context = x[..., -context_size:]
         self._last_sr = sr
         self._last_batch_size = batch_size
 
-        out = torch.from_numpy(out)
+        out = torch.tensor(out)
         return out
 
-    def audio_forward(self, x, sr: int):
+    def audio_forward(self, x, sr: int, num_samples: int = 512):
         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)
@@ -109,7 +94,7 @@ class VoiceActivityDetection():
         return stacked.cpu()
 
     @staticmethod
-    def download(model_url="https://github.com/snakers4/silero-vad/raw/v5.0/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
@@ -153,5 +138,5 @@ class VoiceActivityDetector:
             bool: True if the speech probability exceeds the threshold, indicating the presence of voice activity;
                   False otherwise.
         """
-        speech_probs = self.model.audio_forward(torch.from_numpy(audio_frame.copy()), self.frame_rate)[0]
-        return torch.any(speech_probs > self.threshold).item()
+        speech_prob = self.model(torch.from_numpy(audio_frame), self.frame_rate).item()
+        return speech_prob > self.threshold