add eos to faster_whisper server

Signed-off-by: makaveli10 <vineet.suryan@collabora.com>
2024-07-11 07:01:34 -04:00
19 changed files with 651 additions and 1594 deletions
@@ -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: |
@@ -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
+  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
 )
 ```
 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
+  # Build tiny.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
  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int8   # int8 weight only quantization
  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int4   # int4 weight only quantization
-  # Run server with 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"
                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int8"
                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int4"
  ```
 - CPU
@@ -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
 bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int8   # int8 weight only quantization
 bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int4   # int4 weight only quantization
 # convert small multilingual model
 bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small
@@ -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
 ```
@@ -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.9.0 --extra-index-url https://pypi.nvidia.com
-RUN pip3 install --no-cache-dir -U tensorrt_llm==0.15.0.dev2024111200 --extra-index-url https://pypi.nvidia.com
+
 WORKDIR /app
-RUN git clone https://github.com/NVIDIA/TensorRT-LLM.git && cd TensorRT-LLM && \
+
-    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 .
@@ -1,4 +1,4 @@
 PyAudio
-av
+ffmpeg-python
 scipy
 websocket-client
@@ -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
@@ -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 output_dir="whisper_${model_name//./_}"
    local checkpoint_dir="whisper_${sanitized_model_name}_weights_${weight_only_precision}"
    local output_dir="whisper_${sanitized_model_name}_${weight_only_precision}"
    echo "$output_dir"
-    echo "Converting model weights for $model_name..."
+    echo "Running build script for $model_name with output directory $output_dir"
-    python3 convert_checkpoint.py \
+    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"
-        $( [[ "$weight_only_precision" == "int8" || "$weight_only_precision" == "int4" ]] && echo "--use_weight_only --weight_only_precision $weight_only_precision" ) \
+    echo "Whisper $model_name TensorRT engine built."
        --output_dir "$checkpoint_dir" --model_name "$model_name"
    echo "Building encoder for $model_name..."
    trtllm-build \
        --checkpoint_dir "${checkpoint_dir}/encoder" \
        --output_dir "${output_dir}/encoder" \
        --moe_plugin disable \
        --enable_xqa disable \
        --max_batch_size "$max_batch_size" \
        --gemm_plugin disable \
        --bert_attention_plugin "$inference_precision" \
        --max_input_len 3000 \
        --max_seq_len 3000
    echo "Building decoder for $model_name..."
    trtllm-build \
        --checkpoint_dir "${checkpoint_dir}/decoder" \
        --output_dir "${output_dir}/decoder" \
        --moe_plugin disable \
        --enable_xqa disable \
        --max_beam_width "$max_beam_width" \
        --max_batch_size "$max_batch_size" \
        --max_seq_len 200 \
        --max_input_len 14 \
        --max_encoder_input_len 3000 \
        --gemm_plugin "$inference_precision" \
        --bert_attention_plugin "$inference_precision" \
        --gpt_attention_plugin "$inference_precision"
    echo "TensorRT LLM engine built for $model_name."
    echo "========================================="
    echo "Model is located at: $(pwd)/$output_dir"
 }
@@ -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"
@@ -1,3 +1,3 @@
 #! /bin/bash
-apt-get install portaudio19-dev wget -y
+apt-get install portaudio19-dev ffmpeg wget -y
@@ -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"
 )
@@ -48,9 +48,7 @@ class TestClientCallbacks(BaseTestCase):
            "language": self.client.language,
            "task": self.client.task,
            "model": self.client.model,
-            "use_vad": True,
+            "use_vad": True
            "max_clients": 4,
            "max_connection_time": 600,
        })
        self.client.on_open(self.mock_ws_app)
        self.mock_ws_app.send.assert_called_with(expected_message)
@@ -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": 0, "end": 1, "text": "Test transcript"},
-                {"start": 1, "end": 2, "text": "Test transcript 2", "completed": True},
+                {"start": 1, "end": 2, "text": "Test transcript 2"},
-                {"start": 2, "end": 3, "text": "Test transcript 3", "completed": True}
+                {"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):
@@ -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)
+        wer = self.metric.compute(
-        self.assertLess(wer_score, 0.05)
+            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()
@@ -1 +1 @@
-__version__ = "0.6.3"
+__version__ = "0.5.0"
@@ -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,
+        use_vad=True
        log_transcription=True,
        max_clients=4,
        max_connection_time=600,
    ):
        """
        Initializes a Client instance for audio recording and streaming to a server.
@@ -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.
-            # Truncate to last 3 entries for brevity.
+        text = text[-3:]
-            text = text[-3:]
+        utils.clear_screen()
-            utils.clear_screen()
+        utils.print_transcript(text)
            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,
+                    "use_vad": self.use_vad
                    "max_clients": self.max_clients,
                    "max_connection_time": self.max_connection_time,
                }
            )
        )
@@ -265,9 +250,7 @@ class Client:
        """
        if self.server_backend == "faster_whisper":
-            if not self.transcript and self.last_segment is not None:
+            if (self.last_segment):
                self.transcript.append(self.last_segment)
            elif self.last_segment and self.transcript[-1]["text"] != self.last_segment["text"]:
                self.transcript.append(self.last_segment)
            utils.create_srt_file(self.transcript, output_path)
@@ -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:
+                    self.stream.write(data)
                        time.sleep(chunk_duration)
                    else:
                        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...")
+        process = self.get_rtsp_ffmpeg_process(rtsp_url)
-        try:
+        self.handle_ffmpeg_process(process, stream_type='RTSP')
            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=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")
+            # Process the stream
-            self.process_av_stream(container, stream_type="HLS", save_file=save_file)
+            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):
+        print(f"[INFO]: {stream_type} stream processing finished.")
        """
        Process an AV container stream and send audio packets to the server.
-        Args:
+    def get_rtsp_ffmpeg_process(self, rtsp_url):
-            container (av.container.InputContainer): The input container to process.
+        return (
-            stream_type (str): The type of stream being processed ("RTSP" or "HLS").
+            ffmpeg
-            save_file (str, optional): Local path to save the stream. Default is None.
+            .input(rtsp_url, threads=0)
-        """
+            .output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
-        audio_stream = next((s for s in container.streams if s.type == "audio"), None)
+            .run_async(pipe_stdout=True, pipe_stderr=True)
-        if not audio_stream:
+        )
            print(f"[ERROR]: No audio stream found in {stream_type} source.")
            return
-        output_container = None
+    def get_hls_ffmpeg_process(self, hls_url, save_file):
-        if save_file:
+        if save_file is None:
-            output_container = av.open(save_file, mode="w")
+            process = (
-            output_audio_stream = output_container.add_stream(codec_name="pcm_s16le", rate=self.rate)
+                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:
+        return process
            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):
        """
@@ -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.
+        translate (bool, optional): Indicates whether translation tasks are required (default is False).
-        model (str, optional): The whisper model to use (e.g., "small", "base"). Default is "small".
+        save_output_recording (bool, optional): Indicates whether to save recording from microphone.
-        use_vad (bool, optional): Whether to enable voice activity detection. Default is True.
+        output_recording_filename (str, optional): File to save the output recording.
-        save_output_recording (bool, optional): Whether to save the microphone recording. Default is False.
+        output_transcription_path (str, optional): File to save the output transcription.
        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.
@@ -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",
+        output_transcription_path="./output.srt"
        log_transcription=True,
        max_clients=4,
        max_connection_time=600,
        mute_audio_playback=False,
    ):
-        self.client = Client(
+        self.client = Client(host, port, lang, translate, model, srt_file_path=output_transcription_path, use_vad=use_vad)
            host, port, lang, translate, model, srt_file_path=output_transcription_path,
            use_vad=use_vad, log_transcription=log_transcription, max_clients=max_clients,
            max_connection_time=max_connection_time
        )
        if save_output_recording and not output_recording_filename.endswith(".wav"):
            raise ValueError(f"Please provide a valid `output_recording_filename`: {output_recording_filename}")
        if not output_transcription_path.endswith(".srt"):
@@ -723,6 +675,5 @@ 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
        )
@@ -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 self.backend.is_faster_whisper():
+            if faster_whisper_custom_model_path is not None and os.path.exists(faster_whisper_custom_model_path):
-                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}")
-                    logging.info(f"Using custom model {faster_whisper_custom_model_path}")
+                options["model"] = faster_whisper_custom_model_path
-                    options["model"] = faster_whisper_custom_model_path
+            client = ServeClientFasterWhisper(
-                client = ServeClientFasterWhisper(
+                websocket,
-                    websocket,
+                language=options["language"],
-                    language=options["language"],
+                task=options["task"],
-                    task=options["task"],
+                client_uid=options["uid"],
-                    client_uid=options["uid"],
+                model=options["model"],
-                    model=options["model"],
+                initial_prompt=options.get("initial_prompt"),
-                    initial_prompt=options.get("initial_prompt"),
+                vad_parameters=options.get("vad_parameters"),
-                    vad_parameters=options.get("vad_parameters"),
+                use_vad=self.use_vad,
-                    use_vad=self.use_vad,
+                single_model=self.single_model,
-                    single_model=self.single_model,
+            )
-                )
+            logging.info("Running faster_whisper backend.")
                logging.info("Running faster_whisper backend.")
        except Exception as e:
            return
        if client is None:
            raise ValueError(f"Backend type {self.backend.value} not recognised or not handled.")
@@ -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)
-            voice_active = self.voice_activity(websocket, frame_np)
+        if voice_active:
-            if voice_active:
+            self.no_voice_activity_chunks = 0
-                self.no_voice_activity_chunks = 0
+            client.set_eos(False)
-                client.set_eos(False)
+        if self.use_vad and not voice_active:
-            if self.use_vad and not voice_active:
+            return True
                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:
+        self.single_model = single_model
                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.")
        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:
-            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
-                duration = self.frames_np.shape[0] / self.RATE
+            self.timestamp_offset = self.frames_offset + duration - 5
                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)
-            samples_take = max(0, (self.timestamp_offset - self.frames_offset) * self.RATE)
+        input_bytes = self.frames_np[int(samples_take):].copy()
            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 + " "})
-        
+        self.timestamp_offset += duration
        with self.lock:
            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",
+            "medium", "medium.en", "large-v2", "large-v3",
            "distil-medium.en", "distil-large-v2", "distil-large-v3",
            "large-v3-turbo", "turbo"
        ]
-
+        if not os.path.exists(model):
-        self.model_size_or_path = 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.vad_parameters = vad_parameters or {"threshold": 0.5}
-        self.no_speech_thresh = 0.45
+        self.no_speech_thresh = 0.35
        self.same_output_threshold = 10
        self.end_time_for_same_output = None
        device = "cuda" if torch.cuda.is_available() else "cpu"
        if device == "cuda":
            major, _ = torch.cuda.get_device_capability(device)
            self.compute_type = "float16" if major >= 7 else "float32"
        else:
            self.compute_type = "int8"
        if self.model_size_or_path is None:
            return
        logging.info(f"Using Device={device} with precision {self.compute_type}")
-        try:
+        if single_model:
-            if single_model:
+            if ServeClientFasterWhisper.SINGLE_MODEL is None:
                if ServeClientFasterWhisper.SINGLE_MODEL is None:
                    self.create_model(device)
                    ServeClientFasterWhisper.SINGLE_MODEL = self.transcriber
                else:
                    self.transcriber = ServeClientFasterWhisper.SINGLE_MODEL
            else:
                self.create_model(device)
-        except Exception as e:
+                ServeClientFasterWhisper.SINGLE_MODEL = self.transcriber
-            logging.error(f"Failed to load model: {e}")
+            else:
-            self.websocket.send(json.dumps({
+                print("Re-using already initialized model.")
-                "uid": self.client_uid,
+                self.transcriber = ServeClientFasterWhisper.SINGLE_MODEL
-                "status": "ERROR",
+        else:
-                "message": f"Failed to load model: {str(self.model_size_or_path)}"
+            self.create_model(device)
            }))
            self.websocket.close()
            return
        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_filter=False,
-            vad_parameters=self.vad_parameters if self.use_vad else None)
+            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):
+        if len(self.text):
-            self.send_transcription_to_client(segments)
+            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:
+            if duration < 0.6:
-                time.sleep(0.1)     # wait for audio chunks to arrive
+                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,
+            'text': text
            'completed': completed
        }
    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_)
+                start, end = self.timestamp_offset + s.start, self.timestamp_offset + min(duration, s.end)
                with self.lock:
                    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(
-                last_segment = self.format_segment(
+                self.timestamp_offset + segments[-1].start,
-                    self.timestamp_offset + segments[-1].start,
+                self.timestamp_offset + min(duration, segments[-1].end),
-                    self.timestamp_offset + min(duration, segments[-1].end),
+                self.current_out
-                    self.current_out,
+            )
                    completed=False
                )
        if self.current_out.strip() == self.prev_out.strip() and self.current_out != '':
            self.same_output_count += 1
            # if we remove the audio because of same output on the nth reptition we might remove the 
            # audio thats not yet transcribed so, capturing the time when it was repeated for the first time
            if self.end_time_for_same_output is None:
                self.end_time_for_same_output = segments[-1].end
            time.sleep(0.1)     # wait for some voice activity just in case there is an unitended pause from the speaker for better punctuations.
        else:
            self.same_output_count = 0
            self.end_time_for_same_output = None
        # if same incomplete segment is seen multiple times then update the offset
        # and append the segment to the list
-        if self.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.transcript.append(self.format_segment(
+                    self.timestamp_offset,
-                        self.timestamp_offset,
+                    self.timestamp_offset + duration,
-                        self.timestamp_offset + min(duration, self.end_time_for_same_output),
+                    self.current_out
-                        self.current_out,
+                ))
                        completed=True
                    ))
            self.current_out = ''
-            offset = min(duration, self.end_time_for_same_output)
+            offset = duration
-            self.same_output_count = 0
+            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
@@ -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:
+    # This launches a subprocess to decode audio while down-mixing
-        num_frames = wav_file.getnframes()
+    # and resampling as necessary.  Requires the ffmpeg CLI in PATH.
-        raw_data = wav_file.readframes(num_frames)
+    # 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)
+    return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
    audio_data = audio_data.astype(np.float32) / 32768.0
    return audio_data
 def load_audio_wav_format(wav_path):
@@ -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 ModelConfig, SamplingConfig
 from tensorrt_llm.runtime import PYTHON_BINDINGS, 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,
+    def get_audio_features(self, mel):
-                           mel,
+
-                           mel_input_lengths,
+        input_lengths = torch.tensor(
-                           encoder_downsampling_factor=2):
+            [mel.shape[2] // 2 for _ in range(mel.shape[0])],
        if isinstance(mel, list):
            longest_mel = max([f.shape[-1] for f in mel])
            mel = [
                torch.nn.functional.pad(f, (0, longest_mel - f.shape[-1]),
                                        mode='constant') for f in mel
            ]
            mel = torch.cat(mel, dim=0).type(
                str_dtype_to_torch("float16")).contiguous()
        bsz, seq_len = mel.shape[0], mel.shape[2]
        position_ids = torch.arange(
            math.ceil(seq_len / encoder_downsampling_factor),
            dtype=torch.int32,
-            device=mel.device).expand(bsz, -1).contiguous()
+            device=mel.device)
-        if self.encoder_config['plugin_config']['remove_input_padding']:
+
            # mel B,D,T -> B,T,D -> BxT, D
            mel = mel.transpose(1, 2)
            mel = remove_tensor_padding(mel, mel_input_lengths)
            position_ids = remove_tensor_padding(
                position_ids, mel_input_lengths // encoder_downsampling_factor)
        inputs = OrderedDict()
-        inputs['input_features'] = mel
+        inputs['x'] = mel
-        inputs['input_lengths'] = mel_input_lengths
+        inputs['input_lengths'] = input_lengths
        inputs['position_ids'] = position_ids
        output_list = [
-            TensorInfo('input_features', str_dtype_to_trt(self.dtype),
+            TensorInfo('x', str_dtype_to_trt(self.dtype), mel.shape),
                       mel.shape),
            TensorInfo('input_lengths', str_dtype_to_trt('int32'),
-                       mel_input_lengths.shape),
+                       input_lengths.shape)
            TensorInfo('position_ids', str_dtype_to_trt('int32'),
                       inputs['position_ids'].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']
+        audio_features = outputs['output']
-        encoder_output_lengths = mel_input_lengths // encoder_downsampling_factor
+        return audio_features
        return encoder_output, encoder_output_lengths
 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_heads=self.decoder_config['num_heads'],
-            num_kv_heads=self.decoder_config['num_attention_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_layers'],
-            num_layers=self.decoder_config['num_hidden_layers'],
+            gpt_attention_plugin=self.decoder_config['gpt_attention_plugin'],
-            gpt_attention_plugin=self.decoder_config['plugin_config']
+            remove_input_padding=self.decoder_config['remove_input_padding'],
-            ['gpt_attention_plugin'],
+            cross_attention=self.decoder_config['cross_attention'],
            remove_input_padding=self.decoder_config['plugin_config']
            ['remove_input_padding'],
            kv_cache_type=KVCacheType.PAGED
            if self.decoder_config['plugin_config']['paged_kv_cache'] == True
            else KVCacheType.CONTINUOUS,
            has_position_embedding=self.
            decoder_config['has_position_embedding'],
-            dtype=self.decoder_config['dtype'],
+            has_token_type_embedding=self.
-            has_token_type_embedding=False,
+            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([
+        cross_attention_mask = torch.ones(
-            batch_size, decoder_max_input_length + max_new_tokens,
+            [encoder_outputs.shape[0], 1,
-            encoder_max_input_length
+             encoder_outputs.shape[1]]).int().cuda()
-        ]).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,
+                                       runtime_mapping,
-                                        debug_mode=False)
+                                       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,
+            num_beams=1):
            max_new_tokens=96):
        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_output = self.encoder.get_audio_features(mel)
        encoder_max_input_length = torch.max(encoder_output_lengths).item()
        output_ids = self.decoder.generate(decoder_input_ids,
                                           encoder_output,
                                           encoder_max_input_length,
                                           encoder_output_lengths,
                                           self.tokenizer.eot,
-                                           max_new_tokens=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
+        predictions = self.process_batch(mel, text_prefix, num_beams)
        mel = mel.repeat(batch_size, 1, 1)
        if padding_strategy == "longest":
            pass
        else:
            mel = torch.nn.functional.pad(mel, (0, 3000 - mel.shape[2]))
        features_input_lengths = torch.full((mel.shape[0], ),
                                             mel.shape[2],
                                             dtype=torch.int32,
                                             device=mel.device)
        predictions = self.process_batch(mel, features_input_lengths, text_prefix, num_beams)
        prediction = predictions[0]
        # remove all special tokens in the prediction
@@ -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):
+def create_srt_file(segments, output_file):
-    with open(resampled_file, 'w', encoding='utf-8') as srt_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)
+    try:
-    stream = next(s for s in container.streams if s.type == 'audio')
+        # 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(
+    resampled_file = f"{file.split('.')[0]}_resampled.wav"
-        format='s16',
+    scipy.io.wavfile.write(resampled_file, sr, np_buffer.astype(np.int16))
        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
@@ -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:
+        self.sample_rates = [8000, 16000]
            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:
@@ -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._h = np.zeros((2, batch_size, 64)).astype('float32')
-        self._context = torch.zeros(0)
+        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
+            out, self._h, self._c = 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.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]
+        speech_prob = self.model(torch.from_numpy(audio_frame), self.frame_rate).item()
-        return torch.any(speech_probs > self.threshold).item()
+        return speech_prob > self.threshold
`@@ -1,3 +1,3 @@`
	`#! /bin/bash`	`#! /bin/bash`

	`apt-get install portaudio19-dev wget -y`	`apt-get install portaudio19-dev ffmpeg wget -y`