Bump version v0.7.1

Add transcription callback parameter to TranscriptionClient #361

.github/workflows/ci.yml

@@ -15,7 +15,7 @@ jobs:
     runs-on: ubuntu-22.04
     strategy:
       matrix:
-        python-version: [3.8, 3.9, '3.10', 3.11, 3.12]
+        python-version: [3.9, '3.10', 3.11, 3.12]
     steps:
       - uses: actions/checkout@v2
       
@@ -25,7 +25,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Cache Python dependencies
-        uses: actions/cache@v2
+        uses: actions/cache@v4
         with:
           path: |
             ~/.cache/pip
@@ -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.9, '3.10', 3.11, 3.12]
 
     steps:
       - uses: actions/checkout@v2
@@ -99,35 +99,6 @@ jobs:
           push: true
           tags: ghcr.io/collabora/whisperlive-cpu:latest
   
-  build-and-push-docker-tensorrt:
-    needs: [run-tests, check-code-format]
-    timeout-minutes: 60
-    runs-on: ubuntu-22.04
-    if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/'))
-    steps:
-      - uses: actions/checkout@v2
-
-      - name: Log in to GitHub Container Registry
-        uses: docker/login-action@v1
-        with:
-          registry: ghcr.io
-          username: ${{ github.repository_owner }}
-          password: ${{ secrets.GHCR_TOKEN }}
-
-      - name: Docker Prune
-        run: docker system prune -af
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v1
-
-      - name: Build and push Docker GPU image
-        uses: docker/build-push-action@v2
-        with:
-          context: .
-          file: docker/Dockerfile.tensorrt
-          push: true
-          tags: ghcr.io/collabora/whisperlive-tensorrt:latest
-
   build-and-push-docker-gpu:
     needs: [run-tests, check-code-format, build-and-push-docker-cpu]
     timeout-minutes: 20
@@ -157,6 +128,35 @@ jobs:
           push: true
           tags: ghcr.io/collabora/whisperlive-gpu:latest
 
+  build-and-push-docker-openvino:
+    needs: [run-tests, check-code-format, build-and-push-docker-cpu]
+    timeout-minutes: 20
+    runs-on: ubuntu-22.04
+    if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/'))
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v1
+        with:
+          registry: ghcr.io
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.GHCR_TOKEN }}
+
+      - name: Docker Prune
+        run: docker system prune -af
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+
+      - name: Build and push Docker GPU image
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          file: docker/Dockerfile.openvino
+          push: true
+          tags: ghcr.io/collabora/whisperlive-openvino:latest
+
   publish-to-pypi:
     needs: [run-tests, check-code-format]
     runs-on: ubuntu-22.04
@@ -164,20 +164,20 @@ jobs:
     steps:
       - uses: actions/checkout@v2
 
-      - name: Set up Python 3.8
+      - name: Set up Python 3.9
         uses: actions/setup-python@v2
         with:
-          python-version: 3.8
+          python-version: 3.9
 
       - name: Cache Python dependencies
-        uses: actions/cache@v2
+        uses: actions/cache@v4
         with:
           path: |
             ~/.cache/pip
             !~/.cache/pip/log
-          key: ubuntu-latest-pip-3.8-${{ hashFiles('requirements/server.txt', 'requirements/client.txt') }}
+          key: ubuntu-latest-pip-3.9-${{ hashFiles('requirements/server.txt', 'requirements/client.txt') }}
           restore-keys: |
-            ubuntu-latest-pip-3.8-
+            ubuntu-latest-pip-3.9-
 
       - name: Install system dependencies
         run: sudo apt-get update && sudo apt-get install -y portaudio19-dev

README.md

@@ -11,6 +11,16 @@ This project is a real-time transcription application that uses the OpenAI Whisp
 to convert speech input into text output. It can be used to transcribe both live audio
 input from microphone and pre-recorded audio files.
 
+- [Installation](#installation)
+- [Getting Started](#getting-started)
+- [Running the Server](#running-the-server)
+- [Running the Client](#running-the-client)
+- [Browser Extensions](#browser-extensions)
+- [Whisper Live Server in Docker](#whisper-live-server-in-docker)
+- [Future Work](#future-work)
+- [Contact](#contact)
+- [Citations](#citations)
+
 ## Installation
 - Install PyAudio
 ```bash
@@ -26,7 +36,7 @@ input from microphone and pre-recorded audio files.
 - Please follow [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md) for setup of [NVIDIA/TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) and for building Whisper-TensorRT engine.
 
 ## Getting Started
-The server supports two backends `faster_whisper` and `tensorrt`. If running `tensorrt` backend follow [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md)
+The server supports 3 backends `faster_whisper`, `tensorrt` and `openvino`. If running `tensorrt` backend follow [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md)
 
 ### Running the Server
 - [Faster Whisper](https://github.com/SYSTRAN/faster-whisper) backend
@@ -53,6 +63,16 @@ python3 run_server.py -p 9090 \
                       -trt /home/TensorRT-LLM/examples/whisper/whisper_small \
                       -m
 ```
+
+- WhisperLive now supports the [OpenVINO](https://github.com/openvinotoolkit/openvino) backend for efficient inference on Intel CPUs, iGPU and dGPUs. Currently, we tested the models uploaded to [huggingface by OpenVINO](https://huggingface.co/OpenVINO?search_models=whisper).
+  - > **Docker Recommended:** Running WhisperLive with OpenVINO inside Docker automatically enables GPU support (iGPU/dGPU) without requiring additional host setup.
+  - > **Native (non-Docker) Use:** If you prefer running outside Docker, ensure the Intel drivers and OpenVINO runtime are installed and properly configured on your system. Refer to the documentation for [installing OpenVINO](https://docs.openvino.ai/2025/get-started/install-openvino.html?PACKAGE=OPENVINO_BASE&VERSION=v_2025_0_0&OP_SYSTEM=LINUX&DISTRIBUTION=PIP#).
+
+```
+python3 run_server.py -p 9090 -b openvino
+```
+
+
 #### Controlling OpenMP Threads
 To control the number of threads used by OpenMP, you can set the `OMP_NUM_THREADS` environment variable. This is useful for managing CPU resources and ensuring consistent performance. If not specified, `OMP_NUM_THREADS` is set to `1` by default. You can change this by using the `--omp_num_threads` argument:
 ```bash
@@ -121,7 +141,7 @@ client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/b
 
 ## Browser Extensions
 - Run the server with your desired backend as shown [here](https://github.com/collabora/WhisperLive?tab=readme-ov-file#running-the-server).
-- Transcribe audio directly from your browser using our Chrome or Firefox extensions. Refer to [Audio-Transcription-Chrome](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Chrome#readme) and [Audio-Transcription-Firefox](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Firefox#readme) for setup instructions.
+- Transcribe audio directly from your browser using our Chrome or Firefox extensions. Refer to [Audio-Transcription-Chrome](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Chrome#readme) and https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md
 
 ## Whisper Live Server in Docker
 - GPU
@@ -130,9 +150,10 @@ client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/b
   docker run -it --gpus all -p 9090:9090 ghcr.io/collabora/whisperlive-gpu:latest
   ```
 
-  - TensorRT. 
+  - TensorRT. Refer to [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md) for setup and more tensorrt backend configurations.
   ```bash
-  docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt
+  docker build . -f docker/Dockerfile.tensorrt -t whisperlive-tensorrt
+  docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it whisperlive-tensorrt
 
   # Build small.en engine
   bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en        # float16
@@ -147,15 +168,19 @@ client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/b
                         --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int4"
   ```
 
+  - OpenVINO
+  ```
+  docker run -it --device=/dev/dri -p 9090:9090 ghcr.io/collabora/whisperlive-openvino
+  ```
+
 - CPU
-```bash
-docker run -it -p 9090:9090 ghcr.io/collabora/whisperlive-cpu:latest
-```
-**Note**: By default we use "small" model size. To build docker image for a different model size, change the size in server.py and then build the docker image.
+  - Faster-whisper
+  ```bash
+  docker run -it -p 9090:9090 ghcr.io/collabora/whisperlive-cpu:latest
+  ```
 
 ## Future Work
 - [ ] Add translation to other languages on top of transcription.
-- [x] TensorRT backend for Whisper.
 
 ## Contact
 

TensorRT_whisper.md

@@ -1,6 +1,6 @@
 # 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.18.2`
 
 ## Installation
 - Install [docker](https://docs.docker.com/engine/install/)
@@ -8,7 +8,8 @@ We have only tested the TensorRT backend in docker so, we recommend docker for a
 
 - 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
+docker build . -f docker/Dockerfile.tensorrt -t whisperlive-tensorrt
+docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it whisperlive-tensorrt
 ```
 
 ## Whisper TensorRT Engine
@@ -36,3 +37,11 @@ python3 run_server.py --port 9090 \
                       --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_float16" \
                       --trt_multilingual
 ```
+
+By default trt_backend uses cpp_session, to use python session pass `--trt_py_session` to run_server.py
+```bash
+python3 run_server.py --port 9090 \
+                      --backend tensorrt \
+                      --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_float16" \
+                      --trt_py_session
+```

docker/Dockerfile.openvino

@@ -0,0 +1,19 @@
+FROM openvino/ubuntu22_runtime:latest
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+USER root
+
+RUN apt update && apt install -y portaudio19-dev python-is-python3 && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+RUN pip install --no-cache-dir -U "pip>=24"
+
+RUN mkdir /app
+WORKDIR /app
+
+COPY requirements/server.txt /app/
+RUN pip install --no-cache-dir -r server.txt && rm server.txt
+
+COPY whisper_live /app/whisper_live
+COPY run_server.py /app
+CMD ["python", "run_server.py", "--backend", "openvino"]

docker/Dockerfile.tensorrt

@@ -1,19 +1,19 @@
-FROM nvidia/cuda:12.4.1-base-ubuntu22.04 AS base
+FROM nvidia/cuda:12.8.1-base-ubuntu22.04 AS base
 
 ARG DEBIAN_FRONTEND=noninteractive
 
 RUN apt-get update && apt-get install -y \
     python3.10 python3-pip openmpi-bin libopenmpi-dev git git-lfs wget \
+    && apt install python-is-python3 \
+    && pip install --upgrade pip setuptools \
     && rm -rf /var/lib/apt/lists/*
 
 FROM base AS devel
-RUN pip3 install --no-cache-dir -U tensorrt_llm==0.15.0.dev2024111200 --extra-index-url https://pypi.nvidia.com
+RUN pip install --no-cache-dir -U tensorrt_llm==0.18.2 --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 ../ && \
-    mv TensorRT-LLM/examples ./TensorRT-LLM-examples && \
-    rm -rf TensorRT-LLM
-
+RUN git clone -b v0.18.2 https://github.com/NVIDIA/TensorRT-LLM.git \
+    && mv TensorRT-LLM/examples ./TensorRT-LLM-examples \
+    && rm -rf TensorRT-LLM
 
 FROM devel AS release
 WORKDIR /app
@@ -25,7 +25,6 @@ RUN apt update && bash setup.sh && rm setup.sh
 
 COPY requirements/server.txt .
 RUN pip install --no-cache-dir -r server.txt && rm server.txt
-RUN pip install pynvml==11.5.0
 COPY whisper_live ./whisper_live
 COPY scripts/build_whisper_tensorrt.sh .
 COPY run_server.py .

requirements/server.txt

@@ -10,4 +10,12 @@ jiwer
 evaluate
 numpy<2
 openai-whisper==20240930
-tokenizers==0.20.3
+tokenizers==0.20.3
+
+# openvino
+librosa
+openvino
+openvino-genai
+openvino-tokenizers
+optimum 
+optimum-intel

run_server.py

@@ -10,7 +10,7 @@ if __name__ == "__main__":
     parser.add_argument('--backend', '-b',
                         type=str,
                         default='faster_whisper',
-                        help='Backends from ["tensorrt", "faster_whisper"]')
+                        help='Backends from ["tensorrt", "faster_whisper", "openvino"]')
     parser.add_argument('--faster_whisper_custom_model_path', '-fw',
                         type=str, default=None,
                         help="Custom Faster Whisper Model")
@@ -21,6 +21,9 @@ if __name__ == "__main__":
     parser.add_argument('--trt_multilingual', '-m',
                         action="store_true",
                         help='Boolean only for TensorRT model. True if multilingual.')
+    parser.add_argument('--trt_py_session',
+                        action="store_true",
+                        help='Boolean only for TensorRT model. Use python session or cpp session, By default uses Cpp.')
     parser.add_argument('--omp_num_threads', '-omp',
                         type=int,
                         default=1,
@@ -46,5 +49,6 @@ if __name__ == "__main__":
         faster_whisper_custom_model_path=args.faster_whisper_custom_model_path,
         whisper_tensorrt_path=args.trt_model_path,
         trt_multilingual=args.trt_multilingual,
+        trt_py_session=args.trt_py_session,
         single_model=not args.no_single_model,
     )

scripts/build_whisper_tensorrt.sh

@@ -54,7 +54,7 @@ download_and_build_model() {
     local inference_precision="float16"
     local weight_only_precision="${2:-float16}"
     local max_beam_width=4
-    local max_batch_size=1
+    local max_batch_size=4
 
     echo "Downloading $model_name..."
     # wget --directory-prefix=assets "$model_url"
@@ -80,7 +80,6 @@ download_and_build_model() {
         --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" \
@@ -92,11 +91,10 @@ download_and_build_model() {
         --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_seq_len 225 \
+        --max_input_len 32 \
         --max_encoder_input_len 3000 \
         --gemm_plugin "$inference_precision" \
         --bert_attention_plugin "$inference_precision" \

setup.py

@@ -47,14 +47,21 @@ setup(
         "torch",
         "torchaudio",
         "websockets",
-        "onnxruntime==1.16.0",
+        "onnxruntime==1.17.0",
         "scipy",
         "websocket-client",
         "numba",
         "openai-whisper==20240930",
         "kaldialign",
         "soundfile",
-        "tokenizers==0.20.3"
+        "tokenizers==0.20.3",
+        "librosa",
+        "numpy==1.26.4",
+        "openvino",
+        "openvino-genai",
+        "openvino-tokenizers",
+        "optimum", 
+        "optimum-intel",
     ],
-    python_requires=">=3.8"
+    python_requires=">=3.9"
 )

tests/test_client.py

@@ -51,6 +51,10 @@ class TestClientCallbacks(BaseTestCase):
             "use_vad": True,
             "max_clients": 4,
             "max_connection_time": 600,
+            "send_last_n_segments": 10,
+            "no_speech_thresh": 0.45,
+            "clip_audio": False,
+            "same_output_threshold": 10,
         })
         self.client.on_open(self.mock_ws_app)
         self.mock_ws_app.send.assert_called_with(expected_message)

tests/test_vad.py

@@ -1,6 +1,6 @@
 import unittest
 import numpy as np
-from whisper_live.tensorrt_utils import load_audio
+from whisper_live.transcriber.tensorrt_utils import load_audio
 from whisper_live.vad import VoiceActivityDetector
 
 

whisper_live/__version__.py

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

whisper_live/backend/base.py

@@ -0,0 +1,361 @@
+import json
+import logging
+import threading
+import time
+import numpy as np
+
+
+class ServeClientBase(object):
+    RATE = 16000
+    SERVER_READY = "SERVER_READY"
+    DISCONNECT = "DISCONNECT"
+
+    client_uid: str
+    """A unique identifier for the client."""
+    websocket: object
+    """The WebSocket connection for the client."""
+    send_last_n_segments: int
+    """Number of most recent segments to send to the client."""
+    no_speech_thresh: float
+    """Segments with no speech probability above this threshold will be discarded."""
+    clip_audio: bool
+    """Whether to clip audio with no valid segments."""
+    same_output_threshold: int
+    """Number of repeated outputs before considering it as a valid segment."""
+
+    def __init__(
+        self,
+        client_uid,
+        websocket,
+        send_last_n_segments=10,
+        no_speech_thresh=0.45,
+        clip_audio=False,
+        same_output_threshold=10,
+    ):
+        self.client_uid = client_uid
+        self.websocket = websocket
+        self.send_last_n_segments = send_last_n_segments
+        self.no_speech_thresh = no_speech_thresh
+        self.clip_audio = clip_audio
+        self.same_output_threshold = same_output_threshold
+
+        self.frames = b""
+        self.timestamp_offset = 0.0
+        self.frames_np = None
+        self.frames_offset = 0.0
+        self.text = []
+        self.current_out = ""
+        self.prev_out = ""
+        self.exit = False
+        self.same_output_count = 0
+        self.transcript = []
+        self.end_time_for_same_output = None
+
+        # threading
+        self.lock = threading.Lock()
+
+    def speech_to_text(self):
+        """
+        Process an audio stream in an infinite loop, continuously transcribing the speech.
+
+        This method continuously receives audio frames, performs real-time transcription, and sends
+        transcribed segments to the client via a WebSocket connection.
+
+        If the client's language is not detected, it waits for 30 seconds of audio input to make a language prediction.
+        It utilizes the Whisper ASR model to transcribe the audio, continuously processing and streaming results. Segments
+        are sent to the client in real-time, and a history of segments is maintained to provide context.
+
+        Raises:
+            Exception: If there is an issue with audio processing or WebSocket communication.
+
+        """
+        while True:
+            if self.exit:
+                logging.info("Exiting speech to text thread")
+                break
+
+            if self.frames_np is None:
+                continue
+
+            if self.clip_audio:
+                self.clip_audio_if_no_valid_segment()
+
+            input_bytes, duration = self.get_audio_chunk_for_processing()
+            if duration < 1.0:
+                time.sleep(0.1)     # wait for audio chunks to arrive
+                continue
+            try:
+                input_sample = input_bytes.copy()
+                result = self.transcribe_audio(input_sample)
+
+                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
+                    continue
+                self.handle_transcription_output(result, duration)
+
+            except Exception as e:
+                logging.error(f"[ERROR]: Failed to transcribe audio chunk: {e}")
+                time.sleep(0.01)
+
+    def transcribe_audio(self):
+        raise NotImplementedError
+
+    def handle_transcription_output(self, result, duration):
+        raise NotImplementedError
+    
+    def format_segment(self, start, end, text, completed=False):
+        """
+        Formats a transcription segment with precise start and end times alongside the transcribed text.
+
+        Args:
+            start (float): The start time of the transcription segment in seconds.
+            end (float): The end time of the transcription segment in seconds.
+            text (str): The transcribed text corresponding to the segment.
+
+        Returns:
+            dict: A dictionary representing the formatted transcription segment, including
+                'start' and 'end' times as strings with three decimal places and the 'text'
+                of the transcription.
+        """
+        return {
+            'start': "{:.3f}".format(start),
+            'end': "{:.3f}".format(end),
+            'text': text,
+            'completed': completed
+        }
+
+    def add_frames(self, frame_np):
+        """
+        Add audio frames to the ongoing audio stream buffer.
+
+        This method is responsible for maintaining the audio stream buffer, allowing the continuous addition
+        of audio frames as they are received. It also ensures that the buffer does not exceed a specified size
+        to prevent excessive memory usage.
+
+        If the buffer size exceeds a threshold (45 seconds of audio data), it discards the oldest 30 seconds
+        of audio data to maintain a reasonable buffer size. If the buffer is empty, it initializes it with the provided
+        audio frame. The audio stream buffer is used for real-time processing of audio data for transcription.
+
+        Args:
+            frame_np (numpy.ndarray): The audio frame data as a NumPy array.
+
+        """
+        self.lock.acquire()
+        if self.frames_np is not None and self.frames_np.shape[0] > 45*self.RATE:
+            self.frames_offset += 30.0
+            self.frames_np = self.frames_np[int(30*self.RATE):]
+            # check timestamp offset(should be >= self.frame_offset)
+            # this basically means that there is no speech as timestamp offset hasnt updated
+            # and is less than frame_offset
+            if self.timestamp_offset < self.frames_offset:
+                self.timestamp_offset = self.frames_offset
+        if self.frames_np is None:
+            self.frames_np = frame_np.copy()
+        else:
+            self.frames_np = np.concatenate((self.frames_np, frame_np), axis=0)
+        self.lock.release()
+
+    def clip_audio_if_no_valid_segment(self):
+        """
+        Update the timestamp offset based on audio buffer status.
+        Clip audio if the current chunk exceeds 30 seconds, this basically implies that
+        no valid segment for the last 30 seconds from whisper
+        """
+        with self.lock:
+            if self.frames_np[int((self.timestamp_offset - self.frames_offset)*self.RATE):].shape[0] > 25 * self.RATE:
+                duration = self.frames_np.shape[0] / self.RATE
+                self.timestamp_offset = self.frames_offset + duration - 5
+
+    def get_audio_chunk_for_processing(self):
+        """
+        Retrieves the next chunk of audio data for processing based on the current offsets.
+
+        Calculates which part of the audio data should be processed next, based on
+        the difference between the current timestamp offset and the frame's offset, scaled by
+        the audio sample rate (RATE). It then returns this chunk of audio data along with its
+        duration in seconds.
+
+        Returns:
+            tuple: A tuple containing:
+                - input_bytes (np.ndarray): The next chunk of audio data to be processed.
+                - duration (float): The duration of the audio chunk in seconds.
+        """
+        with self.lock:
+            samples_take = max(0, (self.timestamp_offset - self.frames_offset) * self.RATE)
+            input_bytes = self.frames_np[int(samples_take):].copy()
+        duration = input_bytes.shape[0] / self.RATE
+        return input_bytes, duration
+
+    def prepare_segments(self, last_segment=None):
+        """
+        Prepares the segments of transcribed text to be sent to the client.
+
+        This method compiles the recent segments of transcribed text, ensuring that only the
+        specified number of the most recent segments are included. It also appends the most
+        recent segment of text if provided (which is considered incomplete because of the possibility
+        of the last word being truncated in the audio chunk).
+
+        Args:
+            last_segment (str, optional): The most recent segment of transcribed text to be added
+                                          to the list of segments. Defaults to None.
+
+        Returns:
+            list: A list of transcribed text segments to be sent to the client.
+        """
+        segments = []
+        if len(self.transcript) >= self.send_last_n_segments:
+            segments = self.transcript[-self.send_last_n_segments:].copy()
+        else:
+            segments = self.transcript.copy()
+        if last_segment is not None:
+            segments = segments + [last_segment]
+        return segments
+
+    def get_audio_chunk_duration(self, input_bytes):
+        """
+        Calculates the duration of the provided audio chunk.
+
+        Args:
+            input_bytes (numpy.ndarray): The audio chunk for which to calculate the duration.
+
+        Returns:
+            float: The duration of the audio chunk in seconds.
+        """
+        return input_bytes.shape[0] / self.RATE
+
+    def send_transcription_to_client(self, segments):
+        """
+        Sends the specified transcription segments to the client over the websocket connection.
+
+        This method formats the transcription segments into a JSON object and attempts to send
+        this object to the client. If an error occurs during the send operation, it logs the error.
+
+        Returns:
+            segments (list): A list of transcription segments to be sent to the client.
+        """
+        try:
+            self.websocket.send(
+                json.dumps({
+                    "uid": self.client_uid,
+                    "segments": segments,
+                })
+            )
+        except Exception as e:
+            logging.error(f"[ERROR]: Sending data to client: {e}")
+
+    def disconnect(self):
+        """
+        Notify the client of disconnection and send a disconnect message.
+
+        This method sends a disconnect message to the client via the WebSocket connection to notify them
+        that the transcription service is disconnecting gracefully.
+
+        """
+        self.websocket.send(json.dumps({
+            "uid": self.client_uid,
+            "message": self.DISCONNECT
+        }))
+
+    def cleanup(self):
+        """
+        Perform cleanup tasks before exiting the transcription service.
+
+        This method performs necessary cleanup tasks, including stopping the transcription thread, marking
+        the exit flag to indicate the transcription thread should exit gracefully, and destroying resources
+        associated with the transcription process.
+
+        """
+        logging.info("Cleaning up.")
+        self.exit = True
+    
+    def get_segment_no_speech_prob(self, segment):
+        return getattr(segment, "no_speech_prob", 0)
+
+    def get_segment_start(self, segment):
+        return getattr(segment, "start", getattr(segment, "start_ts", 0))
+
+    def get_segment_end(self, segment):
+        return getattr(segment, "end", getattr(segment, "end_ts", 0))
+
+    def update_segments(self, segments, duration):
+        """
+        Processes the segments from Whisper and updates the transcript.
+        Uses helper methods to account for differences between backends.
+        
+        Args:
+            segments (list): List of segments returned by the transcriber.
+            duration (float): Duration of the current audio chunk.
+        
+        Returns:
+            dict or None: The last processed segment (if any).
+        """
+        offset = None
+        self.current_out = ''
+        last_segment = None
+
+        # Process complete segments only if there are more than one
+        # and if the last segment's no_speech_prob is below the threshold.
+        if len(segments) > 1 and self.get_segment_no_speech_prob(segments[-1]) <= self.no_speech_thresh:
+            for s in segments[:-1]:
+                text_ = s.text
+                self.text.append(text_)
+                with self.lock:
+                    start = self.timestamp_offset + self.get_segment_start(s)
+                    end = self.timestamp_offset + min(duration, self.get_segment_end(s))
+                if start >= end:
+                    continue
+                if self.get_segment_no_speech_prob(s) > self.no_speech_thresh:
+                    continue
+                self.transcript.append(self.format_segment(start, end, text_, completed=True))
+                offset = min(duration, self.get_segment_end(s))
+
+        # Process the last segment if its no_speech_prob is acceptable.
+        if self.get_segment_no_speech_prob(segments[-1]) <= self.no_speech_thresh:
+            self.current_out += segments[-1].text
+            with self.lock:
+                last_segment = self.format_segment(
+                    self.timestamp_offset + self.get_segment_start(segments[-1]),
+                    self.timestamp_offset + min(duration, self.get_segment_end(segments[-1])),
+                    self.current_out,
+                    completed=False
+                )
+
+        # Handle repeated output logic.
+        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 = self.get_segment_end(segments[-1])
+            time.sleep(0.1)  # wait briefly for any new voice activity
+        else:
+            self.same_output_count = 0
+            self.end_time_for_same_output = None
+
+        # If the same incomplete segment is repeated too many times,
+        # append it to the transcript and update the offset.
+        if self.same_output_count > self.same_output_threshold:
+            if not self.text or self.text[-1].strip().lower() != self.current_out.strip().lower():
+                self.text.append(self.current_out)
+                with self.lock:
+                    self.transcript.append(self.format_segment(
+                        self.timestamp_offset,
+                        self.timestamp_offset + min(duration, self.end_time_for_same_output),
+                        self.current_out,
+                        completed=True
+                    ))
+            self.current_out = ''
+            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
+
+        if offset is not None:
+            with self.lock:
+                self.timestamp_offset += offset
+
+        return last_segment

whisper_live/backend/faster_whisper_backend.py

@@ -0,0 +1,216 @@
+import json
+import logging
+import threading
+import time
+import torch
+
+from whisper_live.transcriber.transcriber_faster_whisper import WhisperModel
+from whisper_live.backend.base import ServeClientBase
+
+
+class ServeClientFasterWhisper(ServeClientBase):
+    SINGLE_MODEL = None
+    SINGLE_MODEL_LOCK = threading.Lock()
+
+    def __init__(
+        self,
+        websocket,
+        task="transcribe",
+        device=None,
+        language=None,
+        client_uid=None,
+        model="small.en",
+        initial_prompt=None,
+        vad_parameters=None,
+        use_vad=True,
+        single_model=False,
+        send_last_n_segments=10,
+        no_speech_thresh=0.45,
+        clip_audio=False,
+        same_output_threshold=10,
+    ):
+        """
+        Initialize a ServeClient instance.
+        The Whisper model is initialized based on the client's language and device availability.
+        The transcription thread is started upon initialization. A "SERVER_READY" message is sent
+        to the client to indicate that the server is ready.
+
+        Args:
+            websocket (WebSocket): The WebSocket connection for the client.
+            task (str, optional): The task type, e.g., "transcribe". Defaults to "transcribe".
+            device (str, optional): The device type for Whisper, "cuda" or "cpu". Defaults to None.
+            language (str, optional): The language for transcription. Defaults to None.
+            client_uid (str, optional): A unique identifier for the client. Defaults to None.
+            model (str, optional): The whisper model size. Defaults to 'small.en'
+            initial_prompt (str, optional): Prompt for whisper inference. Defaults to None.
+            single_model (bool, optional): Whether to instantiate a new model for each client connection. Defaults to False.
+            send_last_n_segments (int, optional): Number of most recent segments to send to the client. Defaults to 10.
+            no_speech_thresh (float, optional): Segments with no speech probability above this threshold will be discarded. Defaults to 0.45.
+            clip_audio (bool, optional): Whether to clip audio with no valid segments. Defaults to False.
+            same_output_threshold (int, optional): Number of repeated outputs before considering it as a valid segment. Defaults to 10.
+
+        """
+        super().__init__(
+            client_uid,
+            websocket,
+            send_last_n_segments,
+            no_speech_thresh,
+            clip_audio,
+            same_output_threshold,
+        )
+        self.model_sizes = [
+            "tiny", "tiny.en", "base", "base.en", "small", "small.en",
+            "medium", "medium.en", "large-v2", "large-v3", "distil-small.en",
+            "distil-medium.en", "distil-large-v2", "distil-large-v3",
+            "large-v3-turbo", "turbo"
+        ]
+
+        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}
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        if device == "cuda":
+            major, _ = torch.cuda.get_device_capability(device)
+            self.compute_type = "float16" if major >= 7 else "float32"
+        else:
+            self.compute_type = "int8"
+
+        if self.model_size_or_path is None:
+            return
+        logging.info(f"Using Device={device} with precision {self.compute_type}")
+    
+        try:
+            if single_model:
+                if ServeClientFasterWhisper.SINGLE_MODEL is None:
+                    self.create_model(device)
+                    ServeClientFasterWhisper.SINGLE_MODEL = self.transcriber
+                else:
+                    self.transcriber = ServeClientFasterWhisper.SINGLE_MODEL
+            else:
+                self.create_model(device)
+        except Exception as e:
+            logging.error(f"Failed to load model: {e}")
+            self.websocket.send(json.dumps({
+                "uid": self.client_uid,
+                "status": "ERROR",
+                "message": f"Failed to load model: {str(self.model_size_or_path)}"
+            }))
+            self.websocket.close()
+            return
+
+        self.use_vad = use_vad
+
+        # threading
+        self.trans_thread = threading.Thread(target=self.speech_to_text)
+        self.trans_thread.start()
+        self.websocket.send(
+            json.dumps(
+                {
+                    "uid": self.client_uid,
+                    "message": self.SERVER_READY,
+                    "backend": "faster_whisper"
+                }
+            )
+        )
+
+    def create_model(self, device):
+        """
+        Instantiates a new model, sets it as the transcriber.
+        """
+        self.transcriber = WhisperModel(
+            self.model_size_or_path,
+            device=device,
+            compute_type=self.compute_type,
+            local_files_only=False,
+        )
+
+    def check_valid_model(self, model_size):
+        """
+        Check if it's a valid whisper model size.
+
+        Args:
+            model_size (str): The name of the model size to check.
+
+        Returns:
+            str: The model size if valid, None otherwise.
+        """
+        if model_size not in self.model_sizes:
+            self.websocket.send(
+                json.dumps(
+                    {
+                        "uid": self.client_uid,
+                        "status": "ERROR",
+                        "message": f"Invalid model size {model_size}. Available choices: {self.model_sizes}"
+                    }
+                )
+            )
+            return None
+        return model_size
+
+    def set_language(self, info):
+        """
+        Updates the language attribute based on the detected language information.
+
+        Args:
+            info (object): An object containing the detected language and its probability. This object
+                        must have at least two attributes: `language`, a string indicating the detected
+                        language, and `language_probability`, a float representing the confidence level
+                        of the language detection.
+        """
+        if info.language_probability > 0.5:
+            self.language = info.language
+            logging.info(f"Detected language {self.language} with probability {info.language_probability}")
+            self.websocket.send(json.dumps(
+                {"uid": self.client_uid, "language": self.language, "language_prob": info.language_probability}))
+
+    def transcribe_audio(self, input_sample):
+        """
+        Transcribes the provided audio sample using the configured transcriber instance.
+
+        If the language has not been set, it updates the session's language based on the transcription
+        information.
+
+        Args:
+            input_sample (np.array): The audio chunk to be transcribed. This should be a NumPy
+                                    array representing the audio data.
+
+        Returns:
+            The transcription result from the transcriber. The exact format of this result
+            depends on the implementation of the `transcriber.transcribe` method but typically
+            includes the transcribed text.
+        """
+        if ServeClientFasterWhisper.SINGLE_MODEL:
+            ServeClientFasterWhisper.SINGLE_MODEL_LOCK.acquire()
+        result, info = self.transcriber.transcribe(
+            input_sample,
+            initial_prompt=self.initial_prompt,
+            language=self.language,
+            task=self.task,
+            vad_filter=self.use_vad,
+            vad_parameters=self.vad_parameters if self.use_vad else None)
+        if ServeClientFasterWhisper.SINGLE_MODEL:
+            ServeClientFasterWhisper.SINGLE_MODEL_LOCK.release()
+
+        if self.language is None and info is not None:
+            self.set_language(info)
+        return result
+
+    def handle_transcription_output(self, result, duration):
+        """
+        Handle the transcription output, updating the transcript and sending data to the client.
+
+        Args:
+            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)
+
+        if len(segments):
+            self.send_transcription_to_client(segments)

whisper_live/backend/openvino_backend.py

@@ -0,0 +1,148 @@
+import json
+import logging
+import threading
+import time
+
+from openvino import Core
+from whisper_live.backend.base import ServeClientBase
+from whisper_live.transcriber.transcriber_openvino import WhisperOpenVINO
+
+
+class ServeClientOpenVINO(ServeClientBase):
+    SINGLE_MODEL = None
+    SINGLE_MODEL_LOCK = threading.Lock()
+
+    def __init__(
+        self,
+        websocket,
+        task="transcribe",
+        device=None,
+        language=None,
+        client_uid=None,
+        model="small.en",
+        initial_prompt=None,
+        vad_parameters=None,
+        use_vad=True,
+        single_model=False,
+        send_last_n_segments=10,
+        no_speech_thresh=0.45,
+        clip_audio=False,
+        same_output_threshold=10,
+    ):
+        """
+        Initialize a ServeClient instance.
+        The Whisper model is initialized based on the client's language and device availability.
+        The transcription thread is started upon initialization. A "SERVER_READY" message is sent
+        to the client to indicate that the server is ready.
+
+        Args:
+            websocket (WebSocket): The WebSocket connection for the client.
+            task (str, optional): The task type, e.g., "transcribe." Defaults to "transcribe".
+            device (str, optional): The device type for Whisper, "cuda" or "cpu". Defaults to None.
+            language (str, optional): The language for transcription. Defaults to None.
+            client_uid (str, optional): A unique identifier for the client. Defaults to None.
+            model (str, optional): Huggingface model_id for a valid OpenVINO model.
+            initial_prompt (str, optional): Prompt for whisper inference. Defaults to None.
+            single_model (bool, optional): Whether to instantiate a new model for each client connection. Defaults to False.
+            send_last_n_segments (int, optional): Number of most recent segments to send to the client. Defaults to 10.
+            no_speech_thresh (float, optional): Segments with no speech probability above this threshold will be discarded. Defaults to 0.45.
+            clip_audio (bool, optional): Whether to clip audio with no valid segments. Defaults to False.
+            same_output_threshold (int, optional): Number of repeated outputs before considering it as a valid segment. Defaults to 10.
+        """
+        super().__init__(
+            client_uid,
+            websocket,
+            send_last_n_segments,
+            no_speech_thresh,
+            clip_audio,
+            same_output_threshold,
+        )
+        self.language = "en" if language is None else language
+        if not self.language.startswith("<|"):
+            self.language = f"<|{self.language}|>"
+
+        self.task = "transcribe" if task is None else task
+
+        self.clip_audio = True
+
+        core = Core()
+        available_devices = core.available_devices
+        if 'GPU' in available_devices:
+            selected_device = 'GPU'
+        else:
+            gpu_devices = [d for d in available_devices if d.startswith('GPU')]
+            selected_device = gpu_devices[0] if gpu_devices else 'CPU'
+        self.device = selected_device
+
+
+        if single_model:
+            if ServeClientOpenVINO.SINGLE_MODEL is None:
+                self.create_model(model)
+                ServeClientOpenVINO.SINGLE_MODEL = self.transcriber
+            else:
+                self.transcriber = ServeClientOpenVINO.SINGLE_MODEL
+        else:
+            self.create_model(model)
+
+        # threading
+        self.trans_thread = threading.Thread(target=self.speech_to_text)
+        self.trans_thread.start()
+
+        self.websocket.send(json.dumps({
+            "uid": self.client_uid,
+            "message": self.SERVER_READY,
+            "backend": "openvino"
+        }))
+        logging.info(f"Using OpenVINO device: {self.device}")
+        logging.info(f"Running OpenVINO backend with language: {self.language} and task: {self.task}")
+
+    def create_model(self, model_id):
+        """
+        Instantiates a new model, sets it as the transcriber.
+        """
+        self.transcriber = WhisperOpenVINO(
+            model_id,
+            device=self.device,
+            language=self.language,
+            task=self.task
+        )
+
+    def transcribe_audio(self, input_sample):
+        """
+        Transcribes the provided audio sample using the configured transcriber instance.
+
+        If the language has not been set, it updates the session's language based on the transcription
+        information.
+
+        Args:
+            input_sample (np.array): The audio chunk to be transcribed. This should be a NumPy
+                                    array representing the audio data.
+
+        Returns:
+            The transcription result from the transcriber. The exact format of this result
+            depends on the implementation of the `transcriber.transcribe` method but typically
+            includes the transcribed text.
+        """
+        if ServeClientOpenVINO.SINGLE_MODEL:
+            ServeClientOpenVINO.SINGLE_MODEL_LOCK.acquire()
+        result = self.transcriber.transcribe(input_sample)
+        if ServeClientOpenVINO.SINGLE_MODEL:
+            ServeClientOpenVINO.SINGLE_MODEL_LOCK.release()
+        return result
+
+    def handle_transcription_output(self, result, duration):
+        """
+        Handle the transcription output, updating the transcript and sending data to the client.
+
+        Args:
+            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)
+
+        if len(segments):
+            self.send_transcription_to_client(segments)

whisper_live/backend/trt_backend.py

@@ -0,0 +1,210 @@
+import json
+import logging
+import threading
+import time
+
+from whisper_live.backend.base import ServeClientBase
+from whisper_live.transcriber.transcriber_tensorrt import WhisperTRTLLM
+
+
+class ServeClientTensorRT(ServeClientBase):
+    SINGLE_MODEL = None
+    SINGLE_MODEL_LOCK = threading.Lock()
+
+    def __init__(
+        self,
+        websocket,
+        task="transcribe",
+        multilingual=False,
+        language=None,
+        client_uid=None,
+        model=None,
+        single_model=False,
+        use_py_session=False,
+        max_new_tokens=225,
+        send_last_n_segments=10,
+        no_speech_thresh=0.45,
+        clip_audio=False,
+        same_output_threshold=10,
+    ):
+        """
+        Initialize a ServeClient instance.
+        The Whisper model is initialized based on the client's language and device availability.
+        The transcription thread is started upon initialization. A "SERVER_READY" message is sent
+        to the client to indicate that the server is ready.
+
+        Args:
+            websocket (WebSocket): The WebSocket connection for the client.
+            task (str, optional): The task type, e.g., "transcribe." Defaults to "transcribe".
+            device (str, optional): The device type for Whisper, "cuda" or "cpu". Defaults to None.
+            multilingual (bool, optional): Whether the client supports multilingual transcription. Defaults to False.
+            language (str, optional): The language for transcription. Defaults to None.
+            client_uid (str, optional): A unique identifier for the client. Defaults to None.
+            single_model (bool, optional): Whether to instantiate a new model for each client connection. Defaults to False.
+            use_py_session (bool, optional): Use python session or cpp session. Defaults to Cpp Session.
+            max_new_tokens (int, optional): Max number of tokens to generate.
+            send_last_n_segments (int, optional): Number of most recent segments to send to the client. Defaults to 10.
+            no_speech_thresh (float, optional): Segments with no speech probability above this threshold will be discarded. Defaults to 0.45.
+            clip_audio (bool, optional): Whether to clip audio with no valid segments. Defaults to False.
+            same_output_threshold (int, optional): Number of repeated outputs before considering it as a valid segment. Defaults to 10.
+        """
+        super().__init__(
+            client_uid,
+            websocket,
+            send_last_n_segments,
+            no_speech_thresh,
+            clip_audio,
+            same_output_threshold,
+        )
+
+        self.language = language if multilingual else "en"
+        self.task = task
+        self.eos = False
+        self.max_new_tokens = max_new_tokens
+
+        if single_model:
+            if ServeClientTensorRT.SINGLE_MODEL is None:
+                self.create_model(model, multilingual, use_py_session=use_py_session)
+                ServeClientTensorRT.SINGLE_MODEL = self.transcriber
+            else:
+                self.transcriber = ServeClientTensorRT.SINGLE_MODEL
+        else:
+            self.create_model(model, multilingual, use_py_session=use_py_session)
+
+        # threading
+        self.trans_thread = threading.Thread(target=self.speech_to_text)
+        self.trans_thread.start()
+
+        self.websocket.send(json.dumps({
+            "uid": self.client_uid,
+            "message": self.SERVER_READY,
+            "backend": "tensorrt"
+        }))
+
+    def create_model(self, model, multilingual, warmup=True, use_py_session=False):
+        """
+        Instantiates a new model, sets it as the transcriber and does warmup if desired.
+        """
+        self.transcriber = WhisperTRTLLM(
+            model,
+            assets_dir="assets",
+            device="cuda",
+            is_multilingual=multilingual,
+            language=self.language,
+            task=self.task,
+            use_py_session=use_py_session,
+            max_output_len=self.max_new_tokens,
+        )
+        if warmup:
+            self.warmup()
+
+    def warmup(self, warmup_steps=10):
+        """
+        Warmup TensorRT since first few inferences are slow.
+
+        Args:
+            warmup_steps (int): Number of steps to warm up the model for.
+        """
+        logging.info("[INFO:] Warming up TensorRT engine..")
+        mel, _ = self.transcriber.log_mel_spectrogram("assets/jfk.flac")
+        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.
+
+        Args:
+            last_segment (str): The last segment from the whisper output which is considered to be incomplete because
+                                of the possibility of word being truncated.
+            duration (float): Duration of the transcribed audio chunk.
+        """
+        segments = self.prepare_segments({"text": last_segment})
+        self.send_transcription_to_client(segments)
+        if self.eos:
+            self.update_timestamp_offset(last_segment, duration)
+
+    def transcribe_audio(self, input_bytes):
+        """
+        Transcribe the audio chunk and send the results to the client.
+
+        Args:
+            input_bytes (np.array): The audio chunk to transcribe.
+        """
+        if ServeClientTensorRT.SINGLE_MODEL:
+            ServeClientTensorRT.SINGLE_MODEL_LOCK.acquire()
+        logging.info(f"[WhisperTensorRT:] Processing audio with duration: {input_bytes.shape[0] / self.RATE}")
+        mel, duration = self.transcriber.log_mel_spectrogram(input_bytes)
+        last_segment = self.transcriber.transcribe(
+            mel,
+            text_prefix=f"<|startoftranscript|><|{self.language}|><|{self.task}|><|notimestamps|>",
+        )
+        if ServeClientTensorRT.SINGLE_MODEL:
+            ServeClientTensorRT.SINGLE_MODEL_LOCK.release()
+        if last_segment:
+            self.handle_transcription_output(last_segment, duration)
+
+    def update_timestamp_offset(self, last_segment, duration):
+        """
+        Update timestamp offset and transcript.
+
+        Args:
+            last_segment (str): Last transcribed audio from the whisper model.
+            duration (float): Duration of the last audio chunk.
+        """
+        if not len(self.transcript):
+            self.transcript.append({"text": last_segment + " "})
+        elif self.transcript[-1]["text"].strip() != last_segment:
+            self.transcript.append({"text": last_segment + " "})
+        
+        with self.lock:
+            self.timestamp_offset += duration
+
+    def speech_to_text(self):
+        """
+        Process an audio stream in an infinite loop, continuously transcribing the speech.
+
+        This method continuously receives audio frames, performs real-time transcription, and sends
+        transcribed segments to the client via a WebSocket connection.
+
+        If the client's language is not detected, it waits for 30 seconds of audio input to make a language prediction.
+        It utilizes the Whisper ASR model to transcribe the audio, continuously processing and streaming results. Segments
+        are sent to the client in real-time, and a history of segments is maintained to provide context.
+
+        Raises:
+            Exception: If there is an issue with audio processing or WebSocket communication.
+
+        """
+        while True:
+            if self.exit:
+                logging.info("Exiting speech to text thread")
+                break
+
+            if self.frames_np is None:
+                time.sleep(0.02)    # wait for any audio to arrive
+                continue
+
+            self.clip_audio_if_no_valid_segment()
+
+            input_bytes, duration = self.get_audio_chunk_for_processing()
+            if duration < 0.4:
+                continue
+
+            try:
+                input_sample = input_bytes.copy()
+                logging.info(f"[WhisperTensorRT:] Processing audio with duration: {duration}")
+                self.transcribe_audio(input_sample)
+
+            except Exception as e:
+                logging.error(f"[ERROR]: {e}")

whisper_live/client.py

@@ -30,9 +30,15 @@ class Client:
         model="small",
         srt_file_path="output.srt",
         use_vad=True,
+        use_wss=False,
         log_transcription=True,
         max_clients=4,
         max_connection_time=600,
+        send_last_n_segments=10,
+        no_speech_thresh=0.45,
+        clip_audio=False,
+        same_output_threshold=10,
+        transcription_callback=None,
     ):
         """
         Initializes a Client instance for audio recording and streaming to a server.
@@ -52,6 +58,11 @@ class Client:
             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.
+            send_last_n_segments (int, optional): Number of most recent segments to send to the client. Defaults to 10.
+            no_speech_thresh (float, optional): Segments with no speech probability above this threshold will be discarded. Defaults to 0.45.
+            clip_audio (bool, optional): Whether to clip audio with no valid segments. Defaults to False.
+            same_output_threshold (int, optional): Number of repeated outputs before considering it as a valid segment. Defaults to 10.
+            transcription_callback (callable, optional): A callback function to handle transcription results. Default is None.
         """
         self.recording = False
         self.task = "transcribe"
@@ -64,11 +75,17 @@ class Client:
         self.server_error = False
         self.srt_file_path = srt_file_path
         self.use_vad = use_vad
+        self.use_wss = use_wss
         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
+        self.send_last_n_segments = send_last_n_segments
+        self.no_speech_thresh = no_speech_thresh
+        self.clip_audio = clip_audio
+        self.same_output_threshold = same_output_threshold
+        self.transcription_callback = transcription_callback
 
         if translate:
             self.task = "translate"
@@ -76,7 +93,8 @@ class Client:
         self.audio_bytes = None
 
         if host is not None and port is not None:
-            socket_url = f"ws://{host}:{port}"
+            socket_protocol = 'wss' if self.use_wss else "ws"
+            socket_url = f"{socket_protocol}://{host}:{port}"
             self.client_socket = websocket.WebSocketApp(
                 socket_url,
                 on_open=lambda ws: self.on_open(ws),
@@ -94,7 +112,7 @@ class Client:
 
         # start websocket client in a thread
         self.ws_thread = threading.Thread(target=self.client_socket.run_forever)
-        self.ws_thread.setDaemon(True)
+        self.ws_thread.daemon = True
         self.ws_thread.start()
 
         self.transcript = []
@@ -129,6 +147,14 @@ class Client:
             self.last_response_received = time.time()
             self.last_received_segment = segments[-1]["text"]
 
+        # call the transcription callback if provided
+        if self.transcription_callback and callable(self.transcription_callback):
+            try:
+                self.transcription_callback(" ".join(text), segments) # string, list
+            except Exception as e:
+                print(f"[WARN] transcription_callback raised: {e}")
+            return
+        
         if self.log_transcription:
             # Truncate to last 3 entries for brevity.
             text = text[-3:]
@@ -212,6 +238,10 @@ class Client:
                     "use_vad": self.use_vad,
                     "max_clients": self.max_clients,
                     "max_connection_time": self.max_connection_time,
+                    "send_last_n_segments": self.send_last_n_segments,
+                    "no_speech_thresh": self.no_speech_thresh,
+                    "clip_audio": self.clip_audio,
+                    "same_output_threshold": self.same_output_threshold,
                 }
             )
         )
@@ -682,6 +712,11 @@ class TranscriptionClient(TranscriptionTeeClient):
         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.
+        send_last_n_segments (int, optional): Number of most recent segments to send to the client. Defaults to 10.
+        no_speech_thresh (float, optional): Segments with no speech probability above this threshold will be discarded. Defaults to 0.45.
+        clip_audio (bool, optional): Whether to clip audio with no valid segments. Defaults to False.
+        same_output_threshold (int, optional): Number of repeated outputs before considering it as a valid segment. Defaults to 10.
+        transcription_callback (callable, optional): A callback function to handle transcription results. Default is None.
 
     Attributes:
         client (Client): An instance of the underlying Client class responsible for handling the WebSocket connection.
@@ -701,6 +736,7 @@ class TranscriptionClient(TranscriptionTeeClient):
         translate=False,
         model="small",
         use_vad=True,
+        use_wss=False,
         save_output_recording=False,
         output_recording_filename="./output_recording.wav",
         output_transcription_path="./output.srt",
@@ -708,11 +744,29 @@ class TranscriptionClient(TranscriptionTeeClient):
         max_clients=4,
         max_connection_time=600,
         mute_audio_playback=False,
+        send_last_n_segments=10,
+        no_speech_thresh=0.45,
+        clip_audio=False,
+        same_output_threshold=10,
+        transcription_callback=None,
     ):
         self.client = Client(
-            host, port, lang, translate, model, srt_file_path=output_transcription_path,
-            use_vad=use_vad, log_transcription=log_transcription, max_clients=max_clients,
-            max_connection_time=max_connection_time
+            host,
+            port,
+            lang,
+            translate,
+            model,
+            srt_file_path=output_transcription_path,
+            use_vad=use_vad,
+            use_wss=use_wss,
+            log_transcription=log_transcription,
+            max_clients=max_clients,
+            max_connection_time=max_connection_time,
+            send_last_n_segments=send_last_n_segments,
+            no_speech_thresh=no_speech_thresh,
+            clip_audio=clip_audio,
+            same_output_threshold=same_output_threshold,
+            transcription_callback=transcription_callback,
         )
 
         if save_output_recording and not output_recording_filename.endswith(".wav"):

whisper_live/server.py

@@ -7,16 +7,11 @@ import logging
 from enum import Enum
 from typing import List, Optional
 
-import torch
 import numpy as np
 from websockets.sync.server import serve
 from websockets.exceptions import ConnectionClosed
 from whisper_live.vad import VoiceActivityDetector
-from whisper_live.transcriber import WhisperModel
-try:
-    from whisper_live.transcriber_tensorrt import WhisperTRTLLM
-except Exception:
-    pass
+from whisper_live.backend.base import ServeClientBase
 
 logging.basicConfig(level=logging.INFO)
 
@@ -127,6 +122,7 @@ class ClientManager:
 class BackendType(Enum):
     FASTER_WHISPER = "faster_whisper"
     TENSORRT = "tensorrt"
+    OPENVINO = "openvino"
 
     @staticmethod
     def valid_types() -> List[str]:
@@ -141,6 +137,9 @@ class BackendType(Enum):
 
     def is_tensorrt(self) -> bool:
         return self == BackendType.TENSORRT
+    
+    def is_openvino(self) -> bool:
+        return self == BackendType.OPENVINO
 
 
 class TranscriptionServer:
@@ -154,12 +153,13 @@ class TranscriptionServer:
 
     def initialize_client(
         self, websocket, options, faster_whisper_custom_model_path,
-        whisper_tensorrt_path, trt_multilingual
+        whisper_tensorrt_path, trt_multilingual, trt_py_session=False,
     ):
         client: Optional[ServeClientBase] = None
 
         if self.backend.is_tensorrt():
             try:
+                from whisper_live.backend.trt_backend import ServeClientTensorRT
                 client = ServeClientTensorRT(
                     websocket,
                     multilingual=trt_multilingual,
@@ -168,6 +168,11 @@ class TranscriptionServer:
                     client_uid=options["uid"],
                     model=whisper_tensorrt_path,
                     single_model=self.single_model,
+                    use_py_session=trt_py_session,
+                    send_last_n_segments=options.get("send_last_n_segments", 10),
+                    no_speech_thresh=options.get("no_speech_thresh", 0.45),
+                    clip_audio=options.get("clip_audio", False),
+                    same_output_threshold=options.get("same_output_threshold", 10),
                 )
                 logging.info("Running TensorRT backend.")
             except Exception as e:
@@ -180,9 +185,37 @@ class TranscriptionServer:
                                "Reverting to available backend: 'faster_whisper'"
                 }))
                 self.backend = BackendType.FASTER_WHISPER
+        
+        if self.backend.is_openvino():
+            try:
+                from whisper_live.backend.openvino_backend import ServeClientOpenVINO
+                client = ServeClientOpenVINO(
+                    websocket,
+                    language=options["language"],
+                    task=options["task"],
+                    client_uid=options["uid"],
+                    model=options["model"],
+                    single_model=self.single_model,
+                    send_last_n_segments=options.get("send_last_n_segments", 10),
+                    no_speech_thresh=options.get("no_speech_thresh", 0.45),
+                    clip_audio=options.get("clip_audio", False),
+                    same_output_threshold=options.get("same_output_threshold", 10),
+                )
+                logging.info("Running OpenVINO backend.")
+            except Exception as e:
+                logging.error(f"OpenVINO not supported: {e}")
+                self.backend = BackendType.FASTER_WHISPER
+                self.client_uid = options["uid"]
+                websocket.send(json.dumps({
+                    "uid": self.client_uid,
+                    "status": "WARNING",
+                    "message": "OpenVINO not supported on Server yet. "
+                                "Reverting to available backend: 'faster_whisper'"
+                }))
 
         try:
             if self.backend.is_faster_whisper():
+                from whisper_live.backend.faster_whisper_backend import ServeClientFasterWhisper
                 if faster_whisper_custom_model_path is not None and os.path.exists(faster_whisper_custom_model_path):
                     logging.info(f"Using custom model {faster_whisper_custom_model_path}")
                     options["model"] = faster_whisper_custom_model_path
@@ -196,10 +229,15 @@ class TranscriptionServer:
                     vad_parameters=options.get("vad_parameters"),
                     use_vad=self.use_vad,
                     single_model=self.single_model,
+                    send_last_n_segments=options.get("send_last_n_segments", 10),
+                    no_speech_thresh=options.get("no_speech_thresh", 0.45),
+                    clip_audio=options.get("clip_audio", False),
+                    same_output_threshold=options.get("same_output_threshold", 10),
                 )
 
                 logging.info("Running faster_whisper backend.")
         except Exception as e:
+            logging.error(e)
             return
 
         if client is None:
@@ -223,7 +261,7 @@ class TranscriptionServer:
         return np.frombuffer(frame_data, dtype=np.float32)
 
     def handle_new_connection(self, websocket, faster_whisper_custom_model_path,
-                              whisper_tensorrt_path, trt_multilingual):
+                              whisper_tensorrt_path, trt_multilingual, trt_py_session=False):
         try:
             logging.info("New client connected")
             options = websocket.recv()
@@ -242,7 +280,7 @@ class TranscriptionServer:
             if self.backend.is_tensorrt():
                 self.vad_detector = VoiceActivityDetector(frame_rate=self.RATE)
             self.initialize_client(websocket, options, faster_whisper_custom_model_path,
-                                   whisper_tensorrt_path, trt_multilingual)
+                                   whisper_tensorrt_path, trt_multilingual, trt_py_session=trt_py_session)
             return True
         except json.JSONDecodeError:
             logging.error("Failed to decode JSON from client")
@@ -274,11 +312,12 @@ class TranscriptionServer:
         return True
 
     def recv_audio(self,
-                   websocket,
+                   websocket,   
                    backend: BackendType = BackendType.FASTER_WHISPER,
                    faster_whisper_custom_model_path=None,
                    whisper_tensorrt_path=None,
-                   trt_multilingual=False):
+                   trt_multilingual=False,
+                   trt_py_session=False):
         """
         Receive audio chunks from a client in an infinite loop.
 
@@ -305,7 +344,7 @@ class TranscriptionServer:
         """
         self.backend = backend
         if not self.handle_new_connection(websocket, faster_whisper_custom_model_path,
-                                          whisper_tensorrt_path, trt_multilingual):
+                                          whisper_tensorrt_path, trt_multilingual, trt_py_session=trt_py_session):
             return
 
         try:
@@ -329,6 +368,7 @@ class TranscriptionServer:
             faster_whisper_custom_model_path=None,
             whisper_tensorrt_path=None,
             trt_multilingual=False,
+            trt_py_session=False,
             single_model=False):
         """
         Run the transcription server.
@@ -356,7 +396,8 @@ class TranscriptionServer:
                 backend=BackendType(backend),
                 faster_whisper_custom_model_path=faster_whisper_custom_model_path,
                 whisper_tensorrt_path=whisper_tensorrt_path,
-                trt_multilingual=trt_multilingual
+                trt_multilingual=trt_multilingual,
+                trt_py_session=trt_py_session,
             ),
             host,
             port
@@ -403,732 +444,3 @@ class TranscriptionServer:
         if self.client_manager.get_client(websocket):
             self.client_manager.remove_client(websocket)
 
-
-class ServeClientBase(object):
-    RATE = 16000
-    SERVER_READY = "SERVER_READY"
-    DISCONNECT = "DISCONNECT"
-
-    def __init__(self, client_uid, websocket):
-        self.client_uid = client_uid
-        self.websocket = websocket
-        self.frames = b""
-        self.timestamp_offset = 0.0
-        self.frames_np = None
-        self.frames_offset = 0.0
-        self.text = []
-        self.current_out = ''
-        self.prev_out = ''
-        self.t_start = None
-        self.exit = False
-        self.same_output_count = 0
-        self.show_prev_out_thresh = 5   # if pause(no output from whisper) show previous output for 5 seconds
-        self.add_pause_thresh = 3       # add a blank to segment list as a pause(no speech) for 3 seconds
-        self.transcript = []
-        self.send_last_n_segments = 10
-
-        # text formatting
-        self.pick_previous_segments = 2
-
-        # threading
-        self.lock = threading.Lock()
-
-    def speech_to_text(self):
-        raise NotImplementedError
-
-    def transcribe_audio(self):
-        raise NotImplementedError
-
-    def handle_transcription_output(self):
-        raise NotImplementedError
-
-    def add_frames(self, frame_np):
-        """
-        Add audio frames to the ongoing audio stream buffer.
-
-        This method is responsible for maintaining the audio stream buffer, allowing the continuous addition
-        of audio frames as they are received. It also ensures that the buffer does not exceed a specified size
-        to prevent excessive memory usage.
-
-        If the buffer size exceeds a threshold (45 seconds of audio data), it discards the oldest 30 seconds
-        of audio data to maintain a reasonable buffer size. If the buffer is empty, it initializes it with the provided
-        audio frame. The audio stream buffer is used for real-time processing of audio data for transcription.
-
-        Args:
-            frame_np (numpy.ndarray): The audio frame data as a NumPy array.
-
-        """
-        self.lock.acquire()
-        if self.frames_np is not None and self.frames_np.shape[0] > 45*self.RATE:
-            self.frames_offset += 30.0
-            self.frames_np = self.frames_np[int(30*self.RATE):]
-            # check timestamp offset(should be >= self.frame_offset)
-            # this basically means that there is no speech as timestamp offset hasnt updated
-            # and is less than frame_offset
-            if self.timestamp_offset < self.frames_offset:
-                self.timestamp_offset = self.frames_offset
-        if self.frames_np is None:
-            self.frames_np = frame_np.copy()
-        else:
-            self.frames_np = np.concatenate((self.frames_np, frame_np), axis=0)
-        self.lock.release()
-
-    def clip_audio_if_no_valid_segment(self):
-        """
-        Update the timestamp offset based on audio buffer status.
-        Clip audio if the current chunk exceeds 30 seconds, this basically implies that
-        no valid segment for the last 30 seconds from whisper
-        """
-        with self.lock:
-            if self.frames_np[int((self.timestamp_offset - self.frames_offset)*self.RATE):].shape[0] > 25 * self.RATE:
-                duration = self.frames_np.shape[0] / self.RATE
-                self.timestamp_offset = self.frames_offset + duration - 5
-
-    def get_audio_chunk_for_processing(self):
-        """
-        Retrieves the next chunk of audio data for processing based on the current offsets.
-
-        Calculates which part of the audio data should be processed next, based on
-        the difference between the current timestamp offset and the frame's offset, scaled by
-        the audio sample rate (RATE). It then returns this chunk of audio data along with its
-        duration in seconds.
-
-        Returns:
-            tuple: A tuple containing:
-                - input_bytes (np.ndarray): The next chunk of audio data to be processed.
-                - duration (float): The duration of the audio chunk in seconds.
-        """
-        with self.lock:
-            samples_take = max(0, (self.timestamp_offset - self.frames_offset) * self.RATE)
-            input_bytes = self.frames_np[int(samples_take):].copy()
-        duration = input_bytes.shape[0] / self.RATE
-        return input_bytes, duration
-
-    def prepare_segments(self, last_segment=None):
-        """
-        Prepares the segments of transcribed text to be sent to the client.
-
-        This method compiles the recent segments of transcribed text, ensuring that only the
-        specified number of the most recent segments are included. It also appends the most
-        recent segment of text if provided (which is considered incomplete because of the possibility
-        of the last word being truncated in the audio chunk).
-
-        Args:
-            last_segment (str, optional): The most recent segment of transcribed text to be added
-                                          to the list of segments. Defaults to None.
-
-        Returns:
-            list: A list of transcribed text segments to be sent to the client.
-        """
-        segments = []
-        if len(self.transcript) >= self.send_last_n_segments:
-            segments = self.transcript[-self.send_last_n_segments:].copy()
-        else:
-            segments = self.transcript.copy()
-        if last_segment is not None:
-            segments = segments + [last_segment]
-        return segments
-
-    def get_audio_chunk_duration(self, input_bytes):
-        """
-        Calculates the duration of the provided audio chunk.
-
-        Args:
-            input_bytes (numpy.ndarray): The audio chunk for which to calculate the duration.
-
-        Returns:
-            float: The duration of the audio chunk in seconds.
-        """
-        return input_bytes.shape[0] / self.RATE
-
-    def send_transcription_to_client(self, segments):
-        """
-        Sends the specified transcription segments to the client over the websocket connection.
-
-        This method formats the transcription segments into a JSON object and attempts to send
-        this object to the client. If an error occurs during the send operation, it logs the error.
-
-        Returns:
-            segments (list): A list of transcription segments to be sent to the client.
-        """
-        try:
-            self.websocket.send(
-                json.dumps({
-                    "uid": self.client_uid,
-                    "segments": segments,
-                })
-            )
-        except Exception as e:
-            logging.error(f"[ERROR]: Sending data to client: {e}")
-
-    def disconnect(self):
-        """
-        Notify the client of disconnection and send a disconnect message.
-
-        This method sends a disconnect message to the client via the WebSocket connection to notify them
-        that the transcription service is disconnecting gracefully.
-
-        """
-        self.websocket.send(json.dumps({
-            "uid": self.client_uid,
-            "message": self.DISCONNECT
-        }))
-
-    def cleanup(self):
-        """
-        Perform cleanup tasks before exiting the transcription service.
-
-        This method performs necessary cleanup tasks, including stopping the transcription thread, marking
-        the exit flag to indicate the transcription thread should exit gracefully, and destroying resources
-        associated with the transcription process.
-
-        """
-        logging.info("Cleaning up.")
-        self.exit = True
-
-
-class ServeClientTensorRT(ServeClientBase):
-
-    SINGLE_MODEL = None
-    SINGLE_MODEL_LOCK = threading.Lock()
-
-    def __init__(self, websocket, task="transcribe", multilingual=False, language=None, client_uid=None, model=None, single_model=False):
-        """
-        Initialize a ServeClient instance.
-        The Whisper model is initialized based on the client's language and device availability.
-        The transcription thread is started upon initialization. A "SERVER_READY" message is sent
-        to the client to indicate that the server is ready.
-
-        Args:
-            websocket (WebSocket): The WebSocket connection for the client.
-            task (str, optional): The task type, e.g., "transcribe." Defaults to "transcribe".
-            device (str, optional): The device type for Whisper, "cuda" or "cpu". Defaults to None.
-            multilingual (bool, optional): Whether the client supports multilingual transcription. Defaults to False.
-            language (str, optional): The language for transcription. Defaults to None.
-            client_uid (str, optional): A unique identifier for the client. Defaults to None.
-            single_model (bool, optional): Whether to instantiate a new model for each client connection. Defaults to False.
-
-        """
-        super().__init__(client_uid, websocket)
-        self.language = language if multilingual else "en"
-        self.task = task
-        self.eos = False
-
-        if single_model:
-            if ServeClientTensorRT.SINGLE_MODEL is None:
-                self.create_model(model, multilingual)
-                ServeClientTensorRT.SINGLE_MODEL = self.transcriber
-            else:
-                self.transcriber = ServeClientTensorRT.SINGLE_MODEL
-        else:
-            self.create_model(model, multilingual)
-
-        # threading
-        self.trans_thread = threading.Thread(target=self.speech_to_text)
-        self.trans_thread.start()
-
-        self.websocket.send(json.dumps({
-            "uid": self.client_uid,
-            "message": self.SERVER_READY,
-            "backend": "tensorrt"
-        }))
-
-    def create_model(self, model, multilingual, warmup=True):
-        """
-        Instantiates a new model, sets it as the transcriber and does warmup if desired.
-        """
-        self.transcriber = WhisperTRTLLM(
-            model,
-            assets_dir="assets",
-            device="cuda",
-            is_multilingual=multilingual,
-            language=self.language,
-            task=self.task
-        )
-        if warmup:
-            self.warmup()
-
-    def warmup(self, warmup_steps=10):
-        """
-        Warmup TensorRT since first few inferences are slow.
-
-        Args:
-            warmup_steps (int): Number of steps to warm up the model for.
-        """
-        logging.info("[INFO:] Warming up TensorRT engine..")
-        mel, _ = self.transcriber.log_mel_spectrogram("assets/jfk.flac")
-        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.
-
-        Args:
-            last_segment (str): The last segment from the whisper output which is considered to be incomplete because
-                                of the possibility of word being truncated.
-            duration (float): Duration of the transcribed audio chunk.
-        """
-        segments = self.prepare_segments({"text": last_segment})
-        self.send_transcription_to_client(segments)
-        if self.eos:
-            self.update_timestamp_offset(last_segment, duration)
-
-    def transcribe_audio(self, input_bytes):
-        """
-        Transcribe the audio chunk and send the results to the client.
-
-        Args:
-            input_bytes (np.array): The audio chunk to transcribe.
-        """
-        if ServeClientTensorRT.SINGLE_MODEL:
-            ServeClientTensorRT.SINGLE_MODEL_LOCK.acquire()
-        logging.info(f"[WhisperTensorRT:] Processing audio with duration: {input_bytes.shape[0] / self.RATE}")
-        mel, duration = self.transcriber.log_mel_spectrogram(input_bytes)
-        last_segment = self.transcriber.transcribe(
-            mel,
-            text_prefix=f"<|startoftranscript|><|{self.language}|><|{self.task}|><|notimestamps|>"
-        )
-        if ServeClientTensorRT.SINGLE_MODEL:
-            ServeClientTensorRT.SINGLE_MODEL_LOCK.release()
-        if last_segment:
-            self.handle_transcription_output(last_segment, duration)
-
-    def update_timestamp_offset(self, last_segment, duration):
-        """
-        Update timestamp offset and transcript.
-
-        Args:
-            last_segment (str): Last transcribed audio from the whisper model.
-            duration (float): Duration of the last audio chunk.
-        """
-        if not len(self.transcript):
-            self.transcript.append({"text": last_segment + " "})
-        elif self.transcript[-1]["text"].strip() != last_segment:
-            self.transcript.append({"text": last_segment + " "})
-        
-        with self.lock:
-            self.timestamp_offset += duration
-
-    def speech_to_text(self):
-        """
-        Process an audio stream in an infinite loop, continuously transcribing the speech.
-
-        This method continuously receives audio frames, performs real-time transcription, and sends
-        transcribed segments to the client via a WebSocket connection.
-
-        If the client's language is not detected, it waits for 30 seconds of audio input to make a language prediction.
-        It utilizes the Whisper ASR model to transcribe the audio, continuously processing and streaming results. Segments
-        are sent to the client in real-time, and a history of segments is maintained to provide context.Pauses in speech
-        (no output from Whisper) are handled by showing the previous output for a set duration. A blank segment is added if
-        there is no speech for a specified duration to indicate a pause.
-
-        Raises:
-            Exception: If there is an issue with audio processing or WebSocket communication.
-
-        """
-        while True:
-            if self.exit:
-                logging.info("Exiting speech to text thread")
-                break
-
-            if self.frames_np is None:
-                time.sleep(0.02)    # wait for any audio to arrive
-                continue
-
-            self.clip_audio_if_no_valid_segment()
-
-            input_bytes, duration = self.get_audio_chunk_for_processing()
-            if duration < 0.4:
-                continue
-
-            try:
-                input_sample = input_bytes.copy()
-                logging.info(f"[WhisperTensorRT:] Processing audio with duration: {duration}")
-                self.transcribe_audio(input_sample)
-
-            except Exception as e:
-                logging.error(f"[ERROR]: {e}")
-
-
-class ServeClientFasterWhisper(ServeClientBase):
-
-    SINGLE_MODEL = None
-    SINGLE_MODEL_LOCK = threading.Lock()
-
-    def __init__(self, websocket, task="transcribe", device=None, language=None, client_uid=None, model="small.en",
-                 initial_prompt=None, vad_parameters=None, use_vad=True, single_model=False):
-        """
-        Initialize a ServeClient instance.
-        The Whisper model is initialized based on the client's language and device availability.
-        The transcription thread is started upon initialization. A "SERVER_READY" message is sent
-        to the client to indicate that the server is ready.
-
-        Args:
-            websocket (WebSocket): The WebSocket connection for the client.
-            task (str, optional): The task type, e.g., "transcribe." Defaults to "transcribe".
-            device (str, optional): The device type for Whisper, "cuda" or "cpu". Defaults to None.
-            language (str, optional): The language for transcription. Defaults to None.
-            client_uid (str, optional): A unique identifier for the client. Defaults to None.
-            model (str, optional): The whisper model size. Defaults to 'small.en'
-            initial_prompt (str, optional): Prompt for whisper inference. Defaults to None.
-            single_model (bool, optional): Whether to instantiate a new model for each client connection. Defaults to False.
-        """
-        super().__init__(client_uid, websocket)
-        self.model_sizes = [
-            "tiny", "tiny.en", "base", "base.en", "small", "small.en",
-            "medium", "medium.en", "large-v2", "large-v3", "distil-small.en",
-            "distil-medium.en", "distil-large-v2", "distil-large-v3",
-            "large-v3-turbo", "turbo"
-        ]
-
-        self.model_size_or_path = model
-        self.language = "en" if self.model_size_or_path.endswith("en") else language
-        self.task = task
-        self.initial_prompt = initial_prompt
-        self.vad_parameters = vad_parameters or {"onset": 0.5}
-        self.no_speech_thresh = 0.45
-        self.same_output_threshold = 10
-        self.end_time_for_same_output = None
-
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        if device == "cuda":
-            major, _ = torch.cuda.get_device_capability(device)
-            self.compute_type = "float16" if major >= 7 else "float32"
-        else:
-            self.compute_type = "int8"
-
-        if self.model_size_or_path is None:
-            return
-        logging.info(f"Using Device={device} with precision {self.compute_type}")
-    
-        try:
-            if single_model:
-                if ServeClientFasterWhisper.SINGLE_MODEL is None:
-                    self.create_model(device)
-                    ServeClientFasterWhisper.SINGLE_MODEL = self.transcriber
-                else:
-                    self.transcriber = ServeClientFasterWhisper.SINGLE_MODEL
-            else:
-                self.create_model(device)
-        except Exception as e:
-            logging.error(f"Failed to load model: {e}")
-            self.websocket.send(json.dumps({
-                "uid": self.client_uid,
-                "status": "ERROR",
-                "message": f"Failed to load model: {str(self.model_size_or_path)}"
-            }))
-            self.websocket.close()
-            return
-
-        self.use_vad = use_vad
-
-        # threading
-        self.trans_thread = threading.Thread(target=self.speech_to_text)
-        self.trans_thread.start()
-        self.websocket.send(
-            json.dumps(
-                {
-                    "uid": self.client_uid,
-                    "message": self.SERVER_READY,
-                    "backend": "faster_whisper"
-                }
-            )
-        )
-
-    def create_model(self, device):
-        """
-        Instantiates a new model, sets it as the transcriber.
-        """
-        self.transcriber = WhisperModel(
-            self.model_size_or_path,
-            device=device,
-            compute_type=self.compute_type,
-            local_files_only=False,
-        )
-
-    def check_valid_model(self, model_size):
-        """
-        Check if it's a valid whisper model size.
-
-        Args:
-            model_size (str): The name of the model size to check.
-
-        Returns:
-            str: The model size if valid, None otherwise.
-        """
-        if model_size not in self.model_sizes:
-            self.websocket.send(
-                json.dumps(
-                    {
-                        "uid": self.client_uid,
-                        "status": "ERROR",
-                        "message": f"Invalid model size {model_size}. Available choices: {self.model_sizes}"
-                    }
-                )
-            )
-            return None
-        return model_size
-
-    def set_language(self, info):
-        """
-        Updates the language attribute based on the detected language information.
-
-        Args:
-            info (object): An object containing the detected language and its probability. This object
-                        must have at least two attributes: `language`, a string indicating the detected
-                        language, and `language_probability`, a float representing the confidence level
-                        of the language detection.
-        """
-        if info.language_probability > 0.5:
-            self.language = info.language
-            logging.info(f"Detected language {self.language} with probability {info.language_probability}")
-            self.websocket.send(json.dumps(
-                {"uid": self.client_uid, "language": self.language, "language_prob": info.language_probability}))
-
-    def transcribe_audio(self, input_sample):
-        """
-        Transcribes the provided audio sample using the configured transcriber instance.
-
-        If the language has not been set, it updates the session's language based on the transcription
-        information.
-
-        Args:
-            input_sample (np.array): The audio chunk to be transcribed. This should be a NumPy
-                                    array representing the audio data.
-
-        Returns:
-            The transcription result from the transcriber. The exact format of this result
-            depends on the implementation of the `transcriber.transcribe` method but typically
-            includes the transcribed text.
-        """
-        if ServeClientFasterWhisper.SINGLE_MODEL:
-            ServeClientFasterWhisper.SINGLE_MODEL_LOCK.acquire()
-        result, info = self.transcriber.transcribe(
-            input_sample,
-            initial_prompt=self.initial_prompt,
-            language=self.language,
-            task=self.task,
-            vad_filter=self.use_vad,
-            vad_parameters=self.vad_parameters if self.use_vad else None)
-        if ServeClientFasterWhisper.SINGLE_MODEL:
-            ServeClientFasterWhisper.SINGLE_MODEL_LOCK.release()
-
-        if self.language is None and info is not None:
-            self.set_language(info)
-        return result
-
-    def get_previous_output(self):
-        """
-        Retrieves previously generated transcription outputs if no new transcription is available
-        from the current audio chunks.
-
-        Checks the time since the last transcription output and, if it is within a specified
-        threshold, returns the most recent segments of transcribed text. It also manages
-        adding a pause (blank segment) to indicate a significant gap in speech based on a defined
-        threshold.
-
-        Returns:
-            segments (list): A list of transcription segments. This may include the most recent
-                            transcribed text segments or a blank segment to indicate a pause
-                            in speech.
-        """
-        segments = []
-        if self.t_start is None:
-            self.t_start = time.time()
-        if time.time() - self.t_start < self.show_prev_out_thresh:
-            segments = self.prepare_segments()
-
-        # add a blank if there is no speech for 3 seconds
-        if len(self.text) and self.text[-1] != '':
-            if time.time() - self.t_start > self.add_pause_thresh:
-                self.text.append('')
-        return segments
-
-    def handle_transcription_output(self, result, duration):
-        """
-        Handle the transcription output, updating the transcript and sending data to the client.
-
-        Args:
-            result (str): The result from whisper inference i.e. the list of segments.
-            duration (float): Duration of the transcribed audio chunk.
-        """
-        segments = []
-        if len(result):
-            self.t_start = None
-            last_segment = self.update_segments(result, duration)
-            segments = self.prepare_segments(last_segment)
-        else:
-            # show previous output if there is pause i.e. no output from whisper
-            segments = self.get_previous_output()
-
-        if len(segments):
-            self.send_transcription_to_client(segments)
-
-    def speech_to_text(self):
-        """
-        Process an audio stream in an infinite loop, continuously transcribing the speech.
-
-        This method continuously receives audio frames, performs real-time transcription, and sends
-        transcribed segments to the client via a WebSocket connection.
-
-        If the client's language is not detected, it waits for 30 seconds of audio input to make a language prediction.
-        It utilizes the Whisper ASR model to transcribe the audio, continuously processing and streaming results. Segments
-        are sent to the client in real-time, and a history of segments is maintained to provide context.Pauses in speech
-        (no output from Whisper) are handled by showing the previous output for a set duration. A blank segment is added if
-        there is no speech for a specified duration to indicate a pause.
-
-        Raises:
-            Exception: If there is an issue with audio processing or WebSocket communication.
-
-        """
-        while True:
-            if self.exit:
-                logging.info("Exiting speech to text thread")
-                break
-
-            if self.frames_np is None:
-                continue
-
-            self.clip_audio_if_no_valid_segment()
-
-            input_bytes, duration = self.get_audio_chunk_for_processing()
-            if duration < 1.0:
-                time.sleep(0.1)     # wait for audio chunks to arrive
-                continue
-            try:
-                input_sample = input_bytes.copy()
-                result = self.transcribe_audio(input_sample)
-
-                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
-                    continue
-                self.handle_transcription_output(result, duration)
-
-            except Exception as e:
-                logging.error(f"[ERROR]: Failed to transcribe audio chunk: {e}")
-                time.sleep(0.01)
-
-    def format_segment(self, start, end, text, completed=False):
-        """
-        Formats a transcription segment with precise start and end times alongside the transcribed text.
-
-        Args:
-            start (float): The start time of the transcription segment in seconds.
-            end (float): The end time of the transcription segment in seconds.
-            text (str): The transcribed text corresponding to the segment.
-
-        Returns:
-            dict: A dictionary representing the formatted transcription segment, including
-                'start' and 'end' times as strings with three decimal places and the 'text'
-                of the transcription.
-        """
-        return {
-            'start': "{:.3f}".format(start),
-            'end': "{:.3f}".format(end),
-            'text': text,
-            'completed': completed
-        }
-
-    def update_segments(self, segments, duration):
-        """
-        Processes the segments from whisper. Appends all the segments to the list
-        except for the last segment assuming that it is incomplete.
-
-        Updates the ongoing transcript with transcribed segments, including their start and end times.
-        Complete segments are appended to the transcript in chronological order. Incomplete segments
-        (assumed to be the last one) are processed to identify repeated content. If the same incomplete
-        segment is seen multiple times, it updates the offset and appends the segment to the transcript.
-        A threshold is used to detect repeated content and ensure it is only included once in the transcript.
-        The timestamp offset is updated based on the duration of processed segments. The method returns the
-        last processed segment, allowing it to be sent to the client for real-time updates.
-
-        Args:
-            segments(dict) : dictionary of segments as returned by whisper
-            duration(float): duration of the current chunk
-
-        Returns:
-            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.
-        """
-        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:
-            for i, s in enumerate(segments[:-1]):
-                text_ = s.text
-                self.text.append(text_)
-                with self.lock:
-                    start, end = self.timestamp_offset + s.start, self.timestamp_offset + min(duration, s.end)
-
-                if start >= end:
-                    continue
-                if s.no_speech_prob > self.no_speech_thresh:
-                    continue
-
-                self.transcript.append(self.format_segment(start, end, text_, completed=True))
-                offset = min(duration, s.end)
-
-        # only process the last segment if it satisfies the no_speech_thresh
-        if segments[-1].no_speech_prob <= self.no_speech_thresh:
-            self.current_out += segments[-1].text
-            with self.lock:
-                last_segment = self.format_segment(
-                    self.timestamp_offset + segments[-1].start,
-                    self.timestamp_offset + min(duration, segments[-1].end),
-                    self.current_out,
-                    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 not len(self.text) or self.text[-1].strip().lower() != self.current_out.strip().lower():
-                self.text.append(self.current_out)
-                with self.lock:
-                    self.transcript.append(self.format_segment(
-                        self.timestamp_offset,
-                        self.timestamp_offset + min(duration, self.end_time_for_same_output),
-                        self.current_out,
-                        completed=True
-                    ))
-            self.current_out = ''
-            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
-
-        # update offset
-        if offset is not None:
-            with self.lock:
-                self.timestamp_offset += offset
-
-        return last_segment

whisper_live/transcriber/transcriber_openvino.py

@@ -0,0 +1,23 @@
+import librosa
+import os
+
+import openvino_genai as ov_genai
+import huggingface_hub as hf_hub
+
+
+class WhisperOpenVINO(object):
+    def __init__(self, model_id="OpenVINO/whisper-tiny-fp16-ov", device="CPU", language="en", task="transcribe"):
+        model_path = model_id.split('/')[-1]
+        cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "openvino_whisper_models")
+        os.makedirs(cache_dir, exist_ok=True)
+        model_path = os.path.join(cache_dir, model_path)
+        if not os.path.exists(model_path):
+            hf_hub.snapshot_download(model_id, local_dir=model_path)
+        self.model = ov_genai.WhisperPipeline(str(model_path), device=device)
+        self.language = language
+        self.task = task
+
+    def transcribe(self, input_audio):
+        outputs = self.model.generate(input_audio, return_timestamps=True, language=self.language, task=self.task)
+        outputs = [seg for seg in outputs.chunks]
+        return outputs

whisper_live/transcriber/transcriber_tensorrt.py

@@ -9,7 +9,12 @@ import torch
 import numpy as np
 import torch.nn.functional as F
 from whisper.tokenizer import get_tokenizer
-from whisper_live.tensorrt_utils import (mel_filters, load_audio_wav_format, pad_or_trim, load_audio)
+from whisper_live.transcriber.tensorrt_utils import (
+    mel_filters, 
+    load_audio_wav_format, 
+    pad_or_trim, 
+    load_audio
+)
 
 import tensorrt_llm
 import tensorrt_llm.logger as logger
@@ -18,7 +23,8 @@ from tensorrt_llm._utils import (str_dtype_to_torch, str_dtype_to_trt,
 from tensorrt_llm.bindings import GptJsonConfig, KVCacheType
 from tensorrt_llm.runtime import PYTHON_BINDINGS, ModelConfig, SamplingConfig
 from tensorrt_llm.runtime.session import Session, TensorInfo
-
+if PYTHON_BINDINGS:
+    from tensorrt_llm.runtime import ModelRunnerCpp
 
 SAMPLE_RATE = 16000
 N_FFT = 400
@@ -250,8 +256,17 @@ class WhisperDecoding:
 
 class WhisperTRTLLM(object):
 
-    def __init__(self, engine_dir, assets_dir=None, device=None, is_multilingual=False,
-                 language="en", task="transcribe"):
+    def __init__(self, 
+                 engine_dir,
+                 assets_dir=None,
+                 device=None,
+                 is_multilingual=False,
+                 language="en",
+                 task="transcribe",
+                 use_py_session=False,
+                 num_beams=1,
+                 debug_mode=False,
+                 max_output_len=96):
         world_size = 1
         runtime_rank = tensorrt_llm.mpi_rank()
         runtime_mapping = tensorrt_llm.Mapping(world_size, runtime_rank)
@@ -263,13 +278,6 @@ class WhisperTRTLLM(object):
         self.num_languages = encoder_config['num_languages']
         is_multilingual = (decoder_config['vocab_size'] >= 51865)
 
-        self.encoder = WhisperEncoding(engine_dir)
-        self.decoder = WhisperDecoding(engine_dir,
-                                        runtime_mapping,
-                                        debug_mode=False)
-        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,
@@ -277,7 +285,28 @@ class WhisperTRTLLM(object):
             language=language,
             task=task,
         )
-        self.filters = mel_filters(self.device, self.encoder.n_mels, assets_dir)
+
+        if use_py_session:
+            self.encoder = WhisperEncoding(engine_dir)
+            self.decoder = WhisperDecoding(engine_dir,
+                                        runtime_mapping,
+                                        debug_mode=False)
+        else:
+            json_config = GptJsonConfig.parse_file(engine_dir / 'decoder' /
+                                                   'config.json')
+            assert json_config.model_config.supports_inflight_batching
+            runner_kwargs = dict(engine_dir=engine_dir,
+                                 is_enc_dec=True,
+                                 max_batch_size=1,
+                                 max_input_len=3000,
+                                 max_output_len=max_output_len,
+                                 max_beam_width=num_beams,
+                                 debug_mode=debug_mode,
+                                 kv_cache_free_gpu_memory_fraction=0.9,
+                                 cross_kv_cache_fraction=0.5)
+            self.model_runner_cpp = ModelRunnerCpp.from_dir(**runner_kwargs)
+        self.filters = mel_filters(self.device, self.n_mels, assets_dir)
+        self.use_py_session = use_py_session
 
     def log_mel_spectrogram(
         self,
@@ -350,16 +379,38 @@ class WhisperTRTLLM(object):
         prompt_id = torch.tensor(prompt_id)
         batch_size = mel.shape[0]
         decoder_input_ids = prompt_id.repeat(batch_size, 1)
-
-        encoder_output, encoder_output_lengths = self.encoder.get_audio_features(mel, mel_input_lengths)
-        encoder_max_input_length = torch.max(encoder_output_lengths).item()
-        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,
-                                           num_beams=num_beams)
+        if self.use_py_session:
+            encoder_output, encoder_output_lengths = self.encoder.get_audio_features(mel, mel_input_lengths)
+            encoder_max_input_length = torch.max(encoder_output_lengths).item()
+            output_ids = self.decoder.generate(decoder_input_ids,
+                                            encoder_output,
+                                            encoder_max_input_length,
+                                            encoder_output_lengths,
+                                            self.tokenizer.eot,
+                                            max_new_tokens=max_new_tokens,
+                                            num_beams=num_beams)
+        else:
+            with torch.no_grad():
+                if isinstance(mel, list):
+                    mel = [
+                        m.transpose(1, 2).type(
+                            str_dtype_to_torch("float16")).squeeze(0)
+                        for m in mel
+                    ]
+                else:
+                    mel = mel.transpose(1, 2)
+                outputs = self.model_runner_cpp.generate(
+                    batch_input_ids=decoder_input_ids,
+                    encoder_input_features=mel,
+                    encoder_output_lengths=mel_input_lengths // 2,
+                    max_new_tokens=max_new_tokens,
+                    end_id=self.tokenizer.eot,
+                    pad_id=self.tokenizer.eot,
+                    num_beams=num_beams,
+                    output_sequence_lengths=True,
+                    return_dict=True)
+                torch.cuda.synchronize()
+                output_ids = outputs['output_ids'].cpu().numpy().tolist()
         texts = []
         for i in range(len(output_ids)):
             text = self.tokenizer.decode(output_ids[i][0]).strip()
@@ -374,7 +425,8 @@ class WhisperTRTLLM(object):
             batch_size=1,
             num_beams=1,
             padding_strategy="max",
-            ):
+            max_new_tokens=96,
+        ):
         mel = mel.type(str_dtype_to_torch(dtype))
         mel = mel.unsqueeze(0)
         # repeat the mel spectrogram to match the batch size
@@ -388,7 +440,13 @@ class WhisperTRTLLM(object):
                                              dtype=torch.int32,
                                              device=mel.device)
 
-        predictions = self.process_batch(mel, features_input_lengths, text_prefix, num_beams)
+        predictions = self.process_batch(
+            mel,
+            features_input_lengths,
+            text_prefix,
+            num_beams,
+            max_new_tokens=max_new_tokens
+        )
         prediction = predictions[0]
 
         # remove all special tokens in the prediction