Bump version v0.6.2

Add option to mute audio playback for file input

.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]
+        python-version: [3.8, 3.9, '3.10', 3.11, 3.12]
     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 ffmpeg portaudio19-dev
+        run: sudo apt-get update && sudo apt-get install -y 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]
+        python-version: [3.8, 3.9, '3.10', 3.11, 3.12]
 
     steps:
       - uses: actions/checkout@v2
@@ -180,7 +180,7 @@ jobs:
             ubuntu-latest-pip-3.8-
 
       - name: Install system dependencies
-        run: sudo apt-get update && sudo apt-get install -y ffmpeg portaudio19-dev
+        run: sudo apt-get update && sudo apt-get install -y 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 and ffmpeg
+- Install PyAudio
 ```bash
  bash scripts/setup.sh
 ```
@@ -79,6 +79,7 @@ If you don't want this, set `--no_single_model`.
   - `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
@@ -92,7 +93,8 @@ client = TranscriptionClient(
   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
+  max_connection_time=600,
+  mute_audio_playback=False,                          # Only used for file input, False by Default
 )
 ```
 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.

docker/Dockerfile.tensorrt

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.5.1-runtime-ubuntu22.04 AS base
+FROM nvidia/cuda:12.4.1-base-ubuntu22.04 AS base
 
 ARG DEBIAN_FRONTEND=noninteractive
 
@@ -25,6 +25,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
-ffmpeg-python
+av
 scipy
 websocket-client

requirements/server.txt

@@ -1,11 +1,11 @@
 faster-whisper==1.1.0
 websockets
-onnxruntime==1.16.0
+onnxruntime==1.17.0
 numba
 kaldialign
 soundfile
-ffmpeg-python
 scipy
+av
 jiwer
 evaluate
 numpy<2

scripts/setup.sh

@@ -1,3 +1,3 @@
 #! /bin/bash
 
-apt-get install portaudio19-dev ffmpeg wget -y
+apt-get install portaudio19-dev 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,
@@ -48,7 +48,6 @@ setup(
         "torchaudio",
         "websockets",
         "onnxruntime==1.16.0",
-        "ffmpeg-python",
         "scipy",
         "websocket-client",
         "numba",

whisper_live/__version__.py

@@ -1 +1 @@
-__version__ = "0.6.0"
+__version__ = "0.6.2"

whisper_live/client.py

@@ -10,7 +10,7 @@ import json
 import websocket
 import uuid
 import time
-import ffmpeg
+import av
 import whisper_live.utils as utils
 
 
@@ -46,6 +46,12 @@ 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"
@@ -285,7 +291,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"):
+    def __init__(self, clients, save_output_recording=False, output_recording_filename="./output_recording.wav", mute_audio_playback=False):
         self.clients = clients
         if not self.clients:
             raise Exception("At least one client is required.")
@@ -296,6 +302,7 @@ 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:
@@ -391,6 +398,7 @@ 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)
@@ -399,8 +407,11 @@ class TranscriptionTeeClient:
 
                     audio_array = self.bytes_to_float_array(data)
                     self.multicast_packet(audio_array.tobytes())
-                    self.stream.write(data)
-
+                    if self.mute_audio_playback:
+                        time.sleep(chunk_duration)
+                    else:
+                        self.stream.write(data)
+    
                 wavfile.close()
 
                 for client in self.clients:
@@ -421,84 +432,83 @@ class TranscriptionTeeClient:
 
     def process_rtsp_stream(self, rtsp_url):
         """
-        Connect to an RTSP source, process the audio stream, and send it for trascription.
+        Connect to an RTSP source, process the audio stream, and send it for transcription.
 
         Args:
             rtsp_url (str): The URL of the RTSP stream source.
         """
-        process = self.get_rtsp_ffmpeg_process(rtsp_url)
-        self.handle_ffmpeg_process(process, stream_type='RTSP')
+        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.")
 
-    def process_hls_stream(self, hls_url, save_file):
+    def process_hls_stream(self, hls_url, save_file=None):
         """
         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.
         """
-        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...")
-        stderr_thread = threading.Thread(target=self.consume_stderr, args=(process,))
-        stderr_thread.start()
+        print("[INFO]: Connecting to HLS stream...")
         try:
-            # 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())
-
+            container = av.open(hls_url, format="hls")
+            self.process_av_stream(container, stream_type="HLS", save_file=save_file)
         except Exception as e:
-            print(f"[ERROR]: Failed to connect to {stream_type} stream: {e}")
+            print(f"[ERROR]: Failed to process HLS 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()
-            if process:
-                process.kill()
+        print("[INFO]: HLS stream processing finished.")
 
-        print(f"[INFO]: {stream_type} stream processing finished.")
-
-    def get_rtsp_ffmpeg_process(self, rtsp_url):
-        return (
-            ffmpeg
-            .input(rtsp_url, threads=0)
-            .output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
-            .run_async(pipe_stdout=True, pipe_stderr=True)
-        )
-
-    def get_hls_ffmpeg_process(self, hls_url, save_file):
-        if save_file is None:
-            process = (
-                ffmpeg
-                .input(hls_url, threads=0)
-                .output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
-                .run_async(pipe_stdout=True, pipe_stderr=True)
-            )
-        else:
-            input = ffmpeg.input(hls_url, threads=0)
-            output_file = input.output(save_file, acodec='copy', vcodec='copy').global_args('-loglevel', 'quiet')
-            output_std = input.output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
-            process = (
-                ffmpeg.merge_outputs(output_file, output_std)
-                .run_async(pipe_stdout=True, pipe_stderr=True)
-            )
-
-        return process
-
-    def consume_stderr(self, process):
+    def process_av_stream(self, container, stream_type, save_file=None):
         """
-        Consume and log the stderr output of a process in a separate thread.
+        Process an AV container stream and send audio packets to the server.
 
         Args:
-            process (subprocess.Popen): The process whose stderr output will be logged.
+            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.
         """
-        for line in iter(process.stderr.readline, b""):
-            logging.debug(f'[STDERR]: {line.decode()}')
+        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
+
+        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)
+
+        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()
 
     def save_chunk(self, n_audio_file):
         """
@@ -662,10 +672,16 @@ 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): 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.
+        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.
 
     Attributes:
         client (Client): An instance of the underlying Client class responsible for handling the WebSocket connection.
@@ -691,6 +707,7 @@ class TranscriptionClient(TranscriptionTeeClient):
         log_transcription=True,
         max_clients=4,
         max_connection_time=600,
+        mute_audio_playback=False,
     ):
         self.client = Client(
             host, port, lang, translate, model, srt_file_path=output_transcription_path,
@@ -706,5 +723,6 @@ class TranscriptionClient(TranscriptionTeeClient):
             self,
             [self.client],
             save_output_recording=save_output_recording,
-            output_recording_filename=output_recording_filename
+            output_recording_filename=output_recording_filename,
+            mute_audio_playback=mute_audio_playback
         )

whisper_live/server.py

@@ -718,7 +718,7 @@ class ServeClientTensorRT(ServeClientBase):
         elif self.transcript[-1]["text"].strip() != last_segment:
             self.transcript.append({"text": last_segment + " "})
         
-        with self.lock():
+        with self.lock:
             self.timestamp_offset += duration
 
     def speech_to_text(self):
@@ -800,6 +800,7 @@ class ServeClientFasterWhisper(ServeClientBase):
         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":
@@ -1095,10 +1096,16 @@ class ServeClientFasterWhisper(ServeClientBase):
 
         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.same_output_count > self.same_output_threshold:
@@ -1107,14 +1114,15 @@ class ServeClientFasterWhisper(ServeClientBase):
                 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,
                         completed=True
                     ))
             self.current_out = ''
-            offset = duration
+            offset = min(duration, self.end_time_for_same_output)
             self.same_output_count = 0
             last_segment = None
+            self.end_time_for_same_output = None
         else:
             self.prev_out = self.current_out
 

whisper_live/tensorrt_utils.py

@@ -23,8 +23,12 @@ 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]
 
@@ -35,38 +39,33 @@ CHUNK_LENGTH = 30
 N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE  # 480000 samples in a 30-second chunk
 
 
-def load_audio(file: str, sr: int = SAMPLE_RATE):
+def load_audio(file: str, sr: int = 16000):
     """
-    Open an audio file and read as mono waveform, resampling as necessary
+    Open an audio file, resample it, and read as a mono waveform.
 
     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)
 
-    # 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
+    with wave.open(resampled_file, "rb") as wav_file:
+        num_frames = wav_file.getnframes()
+        raw_data = wav_file.readframes(num_frames)
 
-    return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+        audio_data = np.frombuffer(raw_data, dtype=np.int16)
+
+    audio_data = audio_data.astype(np.float32) / 32768.0
+
+    return audio_data
 
 
 def load_audio_wav_format(wav_path):

whisper_live/utils.py

@@ -1,8 +1,9 @@
 import os
 import textwrap
 import scipy
-import ffmpeg
 import numpy as np
+import av
+from pathlib import Path
 
 
 def clear_screen():
@@ -26,8 +27,8 @@ def format_time(s):
     return f"{hours:02}:{minutes:02}:{seconds:02},{milliseconds:03}"
 
 
-def create_srt_file(segments, output_file):
-    with open(output_file, 'w', encoding='utf-8') as srt_file:
+def create_srt_file(segments, resampled_file):
+    with open(resampled_file, 'w', encoding='utf-8') as srt_file:
         segment_number = 1
         for segment in segments:
             start_time = format_time(float(segment['start']))
@@ -43,9 +44,7 @@ def create_srt_file(segments, output_file):
 
 def resample(file: str, sr: int = 16000):
     """
-    # 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
+    Resample the audio file to 16kHz.
 
     Args:
         file (str): The audio file to open
@@ -54,18 +53,30 @@ def resample(file: str, sr: int = 16000):
     Returns:
         resampled_file (str): The resampled audio file
     """
-    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)
+    container = av.open(file)
+    stream = next(s for s in container.streams if s.type == 'audio')
 
-    resampled_file = f"{file.split('.')[0]}_resampled.wav"
-    scipy.io.wavfile.write(resampled_file, sr, np_buffer.astype(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()
     return resampled_file

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,7 +23,11 @@ class VoiceActivityDetection():
             self.session = onnxruntime.InferenceSession(path, providers=['CUDAExecutionProvider'], sess_options=opts)
 
         self.reset_states()
-        self.sample_rates = [8000, 16000]
+        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):
         if x.dim() == 1:
@@ -34,27 +37,32 @@ 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._h = np.zeros((2, batch_size, 64)).astype('float32')
-        self._c = np.zeros((2, batch_size, 64)).astype('float32')
+        self._state = torch.zeros((2, batch_size, 128)).float()
+        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()
-        return speech_prob > self.threshold
+        speech_probs = self.model.audio_forward(torch.from_numpy(audio_frame.copy()), self.frame_rate)[0]
+        return torch.any(speech_probs > self.threshold).item()