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]
+        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
@@ -35,7 +35,7 @@ jobs:
             ${{ runner.os }}-pip-${{ matrix.python-version }}-
       
       - name: Install system dependencies
-        run: sudo apt-get update && sudo apt-get install -y ffmpeg portaudio19-dev
+        run: sudo apt-get update && sudo apt-get install -y portaudio19-dev
       
       - name: Install Python dependencies
         run: |
@@ -52,7 +52,7 @@ jobs:
     runs-on: ubuntu-22.04
     strategy:
       matrix:
-        python-version: [3.8, 3.9, '3.10', 3.11]
+        python-version: [3.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: 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.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,23 +164,23 @@ 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 ffmpeg portaudio19-dev
+        run: sudo apt-get update && sudo apt-get install -y portaudio19-dev
 
       - name: Install Python dependencies
         run: |

README.md

@@ -11,8 +11,18 @@ 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 and ffmpeg
+- Install PyAudio
 ```bash
  bash scripts/setup.sh
 ```
@@ -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
@@ -36,7 +46,7 @@ python3 run_server.py --port 9090 \
   
 # running with custom model
 python3 run_server.py --port 9090 \
-                      --backend faster_whisper
+                      --backend faster_whisper \
                       -fw "/path/to/custom/faster/whisper/model"
 ```
 
@@ -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
@@ -77,6 +97,10 @@ If you don't want this, set `--no_single_model`.
   - `use_vad`: Whether to use `Voice Activity Detection` on the server.
   - `save_output_recording`: Set to True to save the microphone input as a `.wav` file during live transcription. This option is helpful for recording sessions for later playback or analysis. Defaults to `False`. 
   - `output_recording_filename`: Specifies the `.wav` file path where the microphone input will be saved if `save_output_recording` is set to `True`.
+  - `max_clients`: Specifies the maximum number of clients the server should allow. Defaults to 4.
+  - `max_connection_time`: Maximum connection time for each client in seconds. Defaults to 600.
+  - `mute_audio_playback`: Whether to mute audio playback when transcribing an audio file. Defaults to False.
+
 ```python
 from whisper_live.client import TranscriptionClient
 client = TranscriptionClient(
@@ -84,10 +108,13 @@ client = TranscriptionClient(
   9090,
   lang="en",
   translate=False,
-  model="small",
+  model="small",                                      # also support hf_model => `Systran/faster-whisper-small`
   use_vad=False,
   save_output_recording=True,                         # Only used for microphone input, False by Default
-  output_recording_filename="./output_recording.wav"  # Only used for microphone input
+  output_recording_filename="./output_recording.wav", # Only used for microphone input
+  max_clients=4,
+  max_connection_time=600,
+  mute_audio_playback=False,                          # Only used for file input, False by Default
 )
 ```
 It connects to the server running on localhost at port 9090. Using a multilingual model, language for the transcription will be automatically detected. You can also use the language option to specify the target language for the transcription, in this case, English ("en"). The translate option should be set to `True` if we want to translate from the source language to English and `False` if we want to transcribe in the source language.
@@ -114,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
@@ -123,28 +150,37 @@ 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 tiny.en engine
+  # Build small.en engine
-  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en
+  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en        # float16
+  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int8   # int8 weight only quantization
+  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int4   # int4 weight only quantization
 
-  # Run server with tiny.en
+  # Run server with small.en
   python3 run_server.py --port 9090 \
                         --backend tensorrt \
-                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en"
+                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_float16"
+                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int8"
+                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int4"
+  ```
+
+  - OpenVINO
+  ```
+  docker run -it --device=/dev/dri -p 9090:9090 ghcr.io/collabora/whisperlive-openvino
   ```
 
 - CPU
-```bash
+  - Faster-whisper
-docker run -it -p 9090:9090 ghcr.io/collabora/whisperlive-cpu:latest
+  ```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.
+  ```
 
 ## Future Work
 - [ ] Add translation to other languages on top of transcription.
-- [x] TensorRT backend for Whisper.
 
 ## Contact
 

TensorRT_whisper.md

@@ -1,27 +1,24 @@
 # WhisperLive-TensorRT
 We have only tested the TensorRT backend in docker so, we recommend docker for a smooth TensorRT backend setup.
-**Note**: We use `tensorrt_llm==0.9.0`
+**Note**: We use `tensorrt_llm==0.18.2`
 
 ## Installation
 - Install [docker](https://docs.docker.com/engine/install/)
 - Install [nvidia-container-toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
 
-- Clone this repo.
-```bash
-git clone https://github.com/collabora/WhisperLive.git
-cd WhisperLive
-```
-
 - Run WhisperLive TensorRT in docker
 ```bash
-docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt:latest
+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
 - We build `small.en` and `small` multilingual TensorRT engine as examples below. The script logs the path of the directory with Whisper TensorRT engine. We need that model_path to run the server.
 ```bash
 # convert small.en
-bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en
+bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en        # float16
+bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int8   # int8 weight only quantization
+bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int4   # int4 weight only quantization
 
 # convert small multilingual model
 bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small
@@ -32,11 +29,19 @@ bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small
 # Run English only model
 python3 run_server.py --port 9090 \
                       --backend tensorrt \
-                      --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en"
+                      --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_float16"
 
 # Run Multilingual model
 python3 run_server.py --port 9090 \
                       --backend tensorrt \
-                      --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small" \
+                      --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_float16" \
                       --trt_multilingual
 ```
+
+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,22 @@
-FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
+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 wget \
+    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/*
 
-RUN pip3 install --no-cache-dir -U tensorrt_llm==0.9.0 --extra-index-url https://pypi.nvidia.com
+FROM base AS devel
-
+RUN pip install --no-cache-dir -U tensorrt_llm==0.18.2 --extra-index-url https://pypi.nvidia.com
 WORKDIR /app
+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
 
-RUN git clone -b v0.9.0 --depth 1 https://github.com/NVIDIA/TensorRT-LLM.git && \
+FROM devel AS release
-    mv TensorRT-LLM/examples ./TensorRT-LLM-examples && \
+WORKDIR /app
-    rm -rf TensorRT-LLM
-
 COPY assets/ ./assets
 RUN wget -nc -P assets/ https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz
 
@@ -22,7 +25,6 @@ RUN apt update && bash setup.sh && rm setup.sh
 
 COPY requirements/server.txt .
 RUN pip install --no-cache-dir -r server.txt && rm server.txt
-
 COPY whisper_live ./whisper_live
 COPY scripts/build_whisper_tensorrt.sh .
 COPY run_server.py .

requirements/client.txt

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

requirements/server.txt

@@ -1,12 +1,21 @@
-faster-whisper==1.0.1
+faster-whisper==1.1.0
-torch
 websockets
-onnxruntime==1.16.0
+onnxruntime==1.17.0
 numba
-openai-whisper
 kaldialign
 soundfile
-ffmpeg-python
 scipy
+av
 jiwer
 evaluate
+numpy<2
+openai-whisper==20240930
+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

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

scripts/setup.sh

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

setup.py

@@ -11,7 +11,7 @@ README = (HERE / "README.md").read_text()
 
 # This call to setup() does all the work
 setup(
-    name="whisper-live",
+    name="whisper_live",
     version=__version__,
     description="A nearly-live implementation of OpenAI's Whisper.",
     long_description=README,
@@ -43,18 +43,25 @@ setup(
     ),
     install_requires=[
         "PyAudio",
-        "faster-whisper==1.0.1",
+        "faster-whisper==1.1.0",
         "torch",
         "torchaudio",
         "websockets",
-        "onnxruntime==1.16.0",
+        "onnxruntime==1.17.0",
-        "ffmpeg-python",
         "scipy",
         "websocket-client",
         "numba",
-        "openai-whisper",
+        "openai-whisper==20240930",
         "kaldialign",
         "soundfile",
+        "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

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

tests/test_server.py

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

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.5.0"
+__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

@@ -2,6 +2,7 @@ import os
 import shutil
 import wave
 
+import logging
 import numpy as np
 import pyaudio
 import threading
@@ -9,7 +10,7 @@ import json
 import websocket
 import uuid
 import time
-import ffmpeg
+import av
 import whisper_live.utils as utils
 
 
@@ -28,7 +29,16 @@ class Client:
         translate=False,
         model="small",
         srt_file_path="output.srt",
-        use_vad=True
+        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.
@@ -42,6 +52,17 @@ class Client:
             port (int): The port number for the WebSocket server.
             lang (str, optional): The selected language for transcription. Default is None.
             translate (bool, optional): Specifies if the task is translation. Default is False.
+            model (str, optional): The whisper model to use (e.g., "small", "medium", "large"). Default is "small".
+            srt_file_path (str, optional): The file path to save the output SRT file. Default is "output.srt".
+            use_vad (bool, optional): Whether to enable voice activity detection. Default is True.
+            log_transcription (bool, optional): Whether to log transcription output to the console. Default is True.
+            max_clients (int, optional): Maximum number of client connections allowed. Default is 4.
+            max_connection_time (int, optional): Maximum allowed connection time in seconds. Default is 600.
+            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"
@@ -54,17 +75,26 @@ 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"
 
-        self.timestamp_offset = 0.0
         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),
@@ -82,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 = []
@@ -106,9 +136,9 @@ class Client:
         for i, seg in enumerate(segments):
             if not text or text[-1] != seg["text"]:
                 text.append(seg["text"])
-                if i == len(segments) - 1:
+                if i == len(segments) - 1 and not seg.get("completed", False):
                     self.last_segment = seg
-                elif (self.server_backend == "faster_whisper" and
+                elif (self.server_backend == "faster_whisper" and seg.get("completed", False) and
                       (not self.transcript or
                         float(seg['start']) >= float(self.transcript[-1]['end']))):
                     self.transcript.append(seg)
@@ -117,10 +147,19 @@ class Client:
             self.last_response_received = time.time()
             self.last_received_segment = segments[-1]["text"]
 
-        # Truncate to last 3 entries for brevity.
+        # call the transcription callback if provided
-        text = text[-3:]
+        if self.transcription_callback and callable(self.transcription_callback):
-        utils.clear_screen()
+            try:
-        utils.print_transcript(text)
+                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:]
+            utils.clear_screen()
+            utils.print_transcript(text)
 
     def on_message(self, ws, message):
         """
@@ -196,7 +235,13 @@ class Client:
                     "language": self.language,
                     "task": self.task,
                     "model": self.model,
-                    "use_vad": self.use_vad
+                    "use_vad": self.use_vad,
+                    "max_clients": self.max_clients,
+                    "max_connection_time": self.max_connection_time,
+                    "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,
                 }
             )
         )
@@ -250,7 +295,9 @@ class Client:
 
         """
         if self.server_backend == "faster_whisper":
-            if (self.last_segment):
+            if not self.transcript and self.last_segment is not None:
+                self.transcript.append(self.last_segment)
+            elif self.last_segment and self.transcript[-1]["text"] != self.last_segment["text"]:
                 self.transcript.append(self.last_segment)
             utils.create_srt_file(self.transcript, output_path)
 
@@ -274,7 +321,7 @@ class TranscriptionTeeClient:
     Attributes:
         clients (list): the underlying Client instances responsible for handling WebSocket connections.
     """
-    def __init__(self, clients, save_output_recording=False, output_recording_filename="./output_recording.wav"):
+    def __init__(self, clients, save_output_recording=False, output_recording_filename="./output_recording.wav", mute_audio_playback=False):
         self.clients = clients
         if not self.clients:
             raise Exception("At least one client is required.")
@@ -285,6 +332,7 @@ class TranscriptionTeeClient:
         self.record_seconds = 60000
         self.save_output_recording = save_output_recording
         self.output_recording_filename = output_recording_filename
+        self.mute_audio_playback = mute_audio_playback
         self.frames = b""
         self.p = pyaudio.PyAudio()
         try:
@@ -380,6 +428,7 @@ class TranscriptionTeeClient:
                 output=True,
                 frames_per_buffer=self.chunk,
             )
+            chunk_duration = self.chunk / float(wavfile.getframerate())
             try:
                 while any(client.recording for client in self.clients):
                     data = wavfile.readframes(self.chunk)
@@ -388,7 +437,10 @@ class TranscriptionTeeClient:
 
                     audio_array = self.bytes_to_float_array(data)
                     self.multicast_packet(audio_array.tobytes())
-                    self.stream.write(data)
+                    if self.mute_audio_playback:
+                        time.sleep(chunk_duration)
+                    else:
+                        self.stream.write(data)
     
                 wavfile.close()
 
@@ -410,72 +462,83 @@ class TranscriptionTeeClient:
 
     def process_rtsp_stream(self, rtsp_url):
         """
-        Connect to an RTSP source, process the audio stream, and send it for trascription.
+        Connect to an RTSP source, process the audio stream, and send it for transcription.
 
         Args:
             rtsp_url (str): The URL of the RTSP stream source.
         """
-        process = self.get_rtsp_ffmpeg_process(rtsp_url)
+        print("[INFO]: Connecting to RTSP stream...")
-        self.handle_ffmpeg_process(process, stream_type='RTSP')
+        try:
+            container = av.open(rtsp_url, format="rtsp", options={"rtsp_transport": "tcp"})
+            self.process_av_stream(container, stream_type="RTSP")
+        except Exception as e:
+            print(f"[ERROR]: Failed to process RTSP stream: {e}")
+        finally:
+            for client in self.clients:
+                client.wait_before_disconnect()
+            self.multicast_packet(Client.END_OF_AUDIO.encode('utf-8'), True)
+            self.close_all_clients()
+            self.write_all_clients_srt()
+        print("[INFO]: RTSP stream processing finished.")
 
-    def process_hls_stream(self, hls_url, save_file):
+    def process_hls_stream(self, hls_url, save_file=None):
         """
         Connect to an HLS source, process the audio stream, and send it for transcription.
 
         Args:
             hls_url (str): The URL of the HLS stream source.
-            save_file （str, optional): Local path to save the network stream.
+            save_file (str, optional): Local path to save the network stream.
         """
-        process = self.get_hls_ffmpeg_process(hls_url, save_file)
+        print("[INFO]: Connecting to HLS stream...")
-        self.handle_ffmpeg_process(process, stream_type='HLS')
-
-    def handle_ffmpeg_process(self, process, stream_type):
-        print(f"[INFO]: Connecting to {stream_type} stream...")
         try:
-            # Process the stream
+            container = av.open(hls_url, format="hls")
-            while True:
+            self.process_av_stream(container, stream_type="HLS", save_file=save_file)
-                in_bytes = process.stdout.read(self.chunk * 2)  # 2 bytes per sample
-                if not in_bytes:
-                    break
-                audio_array = self.bytes_to_float_array(in_bytes)
-                self.multicast_packet(audio_array.tobytes())
-
         except Exception as e:
-            print(f"[ERROR]: Failed to connect to {stream_type} stream: {e}")
+            print(f"[ERROR]: Failed to process HLS stream: {e}")
         finally:
+            for client in self.clients:
+                client.wait_before_disconnect()
+            self.multicast_packet(Client.END_OF_AUDIO.encode('utf-8'), True)
             self.close_all_clients()
             self.write_all_clients_srt()
-            if process:
+        print("[INFO]: HLS stream processing finished.")
-                process.kill()
 
-        print(f"[INFO]: {stream_type} stream processing finished.")
+    def process_av_stream(self, container, stream_type, save_file=None):
+        """
+        Process an AV container stream and send audio packets to the server.
 
-    def get_rtsp_ffmpeg_process(self, rtsp_url):
+        Args:
-        return (
+            container (av.container.InputContainer): The input container to process.
-            ffmpeg
+            stream_type (str): The type of stream being processed ("RTSP" or "HLS").
-            .input(rtsp_url, threads=0)
+            save_file (str, optional): Local path to save the stream. Default is None.
-            .output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
+        """
-            .run_async(pipe_stdout=True, pipe_stderr=True)
+        audio_stream = next((s for s in container.streams if s.type == "audio"), None)
-        )
+        if not audio_stream:
+            print(f"[ERROR]: No audio stream found in {stream_type} source.")
+            return
 
-    def get_hls_ffmpeg_process(self, hls_url, save_file):
+        output_container = None
-        if save_file is None:
+        if save_file:
-            process = (
+            output_container = av.open(save_file, mode="w")
-                ffmpeg
+            output_audio_stream = output_container.add_stream(codec_name="pcm_s16le", rate=self.rate)
-                .input(hls_url, threads=0)
-                .output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
-                .run_async(pipe_stdout=True, pipe_stderr=True)
-            )
-        else:
-            input = ffmpeg.input(hls_url, threads=0)
-            output_file = input.output(save_file, acodec='copy', vcodec='copy').global_args('-loglevel', 'quiet')
-            output_std = input.output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
-            process = (
-                ffmpeg.merge_outputs(output_file, output_std)
-                .run_async(pipe_stdout=True, pipe_stderr=True)
-            )
 
-        return process
+        try:
+            for packet in container.demux(audio_stream):
+                for frame in packet.decode():
+                    audio_data = frame.to_ndarray().tobytes()
+                    self.multicast_packet(audio_data)
+
+                    if save_file:
+                        output_container.mux(frame)
+        except Exception as e:
+            print(f"[ERROR]: Error during {stream_type} stream processing: {e}")
+        finally:
+            # Wait for server to send any leftover transcription.
+            time.sleep(5)
+            self.multicast_packet(Client.END_OF_AUDIO.encode('utf-8'), True)
+            if output_container:
+                output_container.close()
+            container.close()
 
     def save_chunk(self, n_audio_file):
         """
@@ -639,10 +702,21 @@ class TranscriptionClient(TranscriptionTeeClient):
         host (str): The hostname or IP address of the server.
         port (int): The port number to connect to on the server.
         lang (str, optional): The primary language for transcription. Default is None, which defaults to English ('en').
-        translate (bool, optional): Indicates whether translation tasks are required (default is False).
+        translate (bool, optional): If True, the task will be translation instead of transcription. Default is False.
-        save_output_recording (bool, optional): Indicates whether to save recording from microphone.
+        model (str, optional): The whisper model to use (e.g., "small", "base"). Default is "small".
-        output_recording_filename (str, optional): File to save the output recording.
+        use_vad (bool, optional): Whether to enable voice activity detection. Default is True.
-        output_transcription_path (str, optional): File to save the output transcription.
+        save_output_recording (bool, optional): Whether to save the microphone recording. Default is False.
+        output_recording_filename (str, optional): Path to save the output recording WAV file. Default is "./output_recording.wav".
+        output_transcription_path (str, optional): File path to save the output transcription (SRT file). Default is "./output.srt".
+        log_transcription (bool, optional): Whether to log transcription output to the console. Default is True.
+        max_clients (int, optional): Maximum number of client connections allowed. Default is 4.
+        max_connection_time (int, optional): Maximum allowed connection time in seconds. Default is 600.
+        mute_audio_playback (bool, optional): If True, mutes audio playback during file playback. Default is False.
+        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.
@@ -662,11 +736,39 @@ 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"
+        output_transcription_path="./output.srt",
+        log_transcription=True,
+        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)
+        self.client = Client(
+            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"):
             raise ValueError(f"Please provide a valid `output_recording_filename`: {output_recording_filename}")
         if not output_transcription_path.endswith(".srt"):
@@ -675,5 +777,6 @@ class TranscriptionClient(TranscriptionTeeClient):
             self,
             [self.client],
             save_output_recording=save_output_recording,
-            output_recording_filename=output_recording_filename
+            output_recording_filename=output_recording_filename,
+            mute_audio_playback=mute_audio_playback
         )

whisper_live/server.py

@@ -4,16 +4,14 @@ import threading
 import json
 import functools
 import logging
-import torch
+from enum import Enum
+from typing import List, Optional
+
 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
+from whisper_live.backend.base import ServeClientBase
-try:
-    from whisper_live.transcriber_tensorrt import WhisperTRTLLM
-except Exception:
-    pass
 
 logging.basicConfig(level=logging.INFO)
 
@@ -121,21 +119,47 @@ class ClientManager:
         return False
 
 
+class BackendType(Enum):
+    FASTER_WHISPER = "faster_whisper"
+    TENSORRT = "tensorrt"
+    OPENVINO = "openvino"
+
+    @staticmethod
+    def valid_types() -> List[str]:
+        return [backend_type.value for backend_type in BackendType]
+
+    @staticmethod
+    def is_valid(backend: str) -> bool:
+        return backend in BackendType.valid_types()
+
+    def is_faster_whisper(self) -> bool:
+        return self == BackendType.FASTER_WHISPER
+
+    def is_tensorrt(self) -> bool:
+        return self == BackendType.TENSORRT
+    
+    def is_openvino(self) -> bool:
+        return self == BackendType.OPENVINO
+
+
 class TranscriptionServer:
     RATE = 16000
 
     def __init__(self):
-        self.client_manager = ClientManager()
+        self.client_manager = None
         self.no_voice_activity_chunks = 0
         self.use_vad = True
         self.single_model = False
 
     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,
     ):
-        if self.backend == "tensorrt":
+        client: Optional[ServeClientBase] = None
+
+        if self.backend.is_tensorrt():
             try:
+                from whisper_live.backend.trt_backend import ServeClientTensorRT
                 client = ServeClientTensorRT(
                     websocket,
                     multilingual=trt_multilingual,
@@ -144,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:
@@ -155,24 +184,64 @@ class TranscriptionServer:
                     "message": "TensorRT-LLM not supported on Server yet. "
                                "Reverting to available backend: 'faster_whisper'"
                 }))
-                self.backend = "faster_whisper"
+                self.backend = BackendType.FASTER_WHISPER
         
-        if self.backend == "faster_whisper":
+        if self.backend.is_openvino():
-            if faster_whisper_custom_model_path is not None and os.path.exists(faster_whisper_custom_model_path):
+            try:
-                logging.info(f"Using custom model {faster_whisper_custom_model_path}")
+                from whisper_live.backend.openvino_backend import ServeClientOpenVINO
-                options["model"] = faster_whisper_custom_model_path
+                client = ServeClientOpenVINO(
-            client = ServeClientFasterWhisper(
+                    websocket,
-                websocket,
+                    language=options["language"],
-                language=options["language"],
+                    task=options["task"],
-                task=options["task"],
+                    client_uid=options["uid"],
-                client_uid=options["uid"],
+                    model=options["model"],
-                model=options["model"],
+                    single_model=self.single_model,
-                initial_prompt=options.get("initial_prompt"),
+                    send_last_n_segments=options.get("send_last_n_segments", 10),
-                vad_parameters=options.get("vad_parameters"),
+                    no_speech_thresh=options.get("no_speech_thresh", 0.45),
-                use_vad=self.use_vad,
+                    clip_audio=options.get("clip_audio", False),
-                single_model=self.single_model,
+                    same_output_threshold=options.get("same_output_threshold", 10),
-            )
+                )
-            logging.info("Running faster_whisper backend.")
+                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
+                client = ServeClientFasterWhisper(
+                    websocket,
+                    language=options["language"],
+                    task=options["task"],
+                    client_uid=options["uid"],
+                    model=options["model"],
+                    initial_prompt=options.get("initial_prompt"),
+                    vad_parameters=options.get("vad_parameters"),
+                    use_vad=self.use_vad,
+                    single_model=self.single_model,
+                    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:
+            raise ValueError(f"Backend type {self.backend.value} not recognised or not handled.")
 
         self.client_manager.add_client(websocket, client)
 
@@ -192,20 +261,26 @@ 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()
             options = json.loads(options)
+
+            if self.client_manager is None:
+                max_clients = options.get('max_clients', 4)
+                max_connection_time = options.get('max_connection_time', 600)
+                self.client_manager = ClientManager(max_clients, max_connection_time)
+
             self.use_vad = options.get('use_vad')
             if self.client_manager.is_server_full(websocket, options):
                 websocket.close()
                 return False  # Indicates that the connection should not continue
 
-            if self.backend == "tensorrt":
+            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")
@@ -221,11 +296,11 @@ class TranscriptionServer:
         frame_np = self.get_audio_from_websocket(websocket)
         client = self.client_manager.get_client(websocket)
         if frame_np is False:
-            if self.backend == "tensorrt":
+            if self.backend.is_tensorrt():
                 client.set_eos(True)
             return False
 
-        if self.backend == "tensorrt":
+        if self.backend.is_tensorrt():
             voice_active = self.voice_activity(websocket, frame_np)
             if voice_active:
                 self.no_voice_activity_chunks = 0
@@ -238,10 +313,11 @@ class TranscriptionServer:
 
     def recv_audio(self,
                    websocket,   
-                   backend="faster_whisper",
+                   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.
 
@@ -268,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:
@@ -292,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.
@@ -311,13 +388,16 @@ class TranscriptionServer:
                 # TODO: load model initially
             else:
                 logging.info("Single model mode currently only works with custom models.")
+        if not BackendType.is_valid(backend):
+            raise ValueError(f"{backend} is not a valid backend type. Choose backend from {BackendType.valid_types()}")
         with serve(
             functools.partial(
                 self.recv_audio,
-                backend=backend,
+                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
@@ -364,690 +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_threshold = 0
-        self.show_prev_out_thresh = 5   # if pause(no output from whisper) show previous output for 5 seconds
-        self.add_pause_thresh = 3       # add a blank to segment list as a pause(no speech) for 3 seconds
-        self.transcript = []
-        self.send_last_n_segments = 10
-
-        # 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
-        """
-        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.
-        """
-        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 + " "})
-        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",
-        ]
-        if not os.path.exists(model):
-            self.model_size_or_path = self.check_valid_model(model)
-        else:
-            self.model_size_or_path = model
-        self.language = "en" if self.model_size_or_path.endswith("en") else language
-        self.task = task
-        self.initial_prompt = initial_prompt
-        self.vad_parameters = vad_parameters or {"threshold": 0.5}
-        self.no_speech_thresh = 0.45
-
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-
-        if self.model_size_or_path is None:
-            return
-
-        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)
-
-        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="int8" if device == "cpu" else "float16",
-            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:
-                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):
-        """
-        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
-        }
-
-    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 = ''
-        # process complete segments
-        if len(segments) > 1:
-            for i, s in enumerate(segments[:-1]):
-                text_ = s.text
-                self.text.append(text_)
-                start, end = self.timestamp_offset + s.start, self.timestamp_offset + min(duration, s.end)
-
-                if start >= end:
-                    continue
-                if s.no_speech_prob > self.no_speech_thresh:
-                    continue
-
-                self.transcript.append(self.format_segment(start, end, text_))
-                offset = min(duration, s.end)
-
-        self.current_out += segments[-1].text
-        last_segment = self.format_segment(
-            self.timestamp_offset + segments[-1].start,
-            self.timestamp_offset + min(duration, segments[-1].end),
-            self.current_out
-        )
-
-        # if same incomplete segment is seen multiple times then update the offset
-        # and append the segment to the list
-        if self.current_out.strip() == self.prev_out.strip() and self.current_out != '':
-            self.same_output_threshold += 1
-        else:
-            self.same_output_threshold = 0
-
-        if self.same_output_threshold > 5:
-            if not len(self.text) or self.text[-1].strip().lower() != self.current_out.strip().lower():
-                self.text.append(self.current_out)
-                self.transcript.append(self.format_segment(
-                    self.timestamp_offset,
-                    self.timestamp_offset + duration,
-                    self.current_out
-                ))
-            self.current_out = ''
-            offset = duration
-            self.same_output_threshold = 0
-            last_segment = None
-        else:
-            self.prev_out = self.current_out
-
-        # update offset
-        if offset is not None:
-            self.timestamp_offset += offset
-
-        return last_segment

whisper_live/transcriber/tensorrt_utils.py

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

whisper_live/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

@@ -0,0 +1,479 @@
+import json
+import re
+import math
+from collections import OrderedDict
+from pathlib import Path
+from typing import Union
+
+import torch
+import numpy as np
+import torch.nn.functional as F
+from whisper.tokenizer import get_tokenizer
+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
+from tensorrt_llm._utils import (str_dtype_to_torch, str_dtype_to_trt,
+                                 trt_dtype_to_torch)
+from tensorrt_llm.bindings import GptJsonConfig, KVCacheType
+from tensorrt_llm.runtime import PYTHON_BINDINGS, ModelConfig, SamplingConfig
+from tensorrt_llm.runtime.session import Session, TensorInfo
+if PYTHON_BINDINGS:
+    from tensorrt_llm.runtime import ModelRunnerCpp
+
+SAMPLE_RATE = 16000
+N_FFT = 400
+HOP_LENGTH = 160
+CHUNK_LENGTH = 30
+N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE  # 480000 samples in a 30-second chunk
+
+def read_config(component, engine_dir):
+    config_path = engine_dir / component / 'config.json'
+    with open(config_path, 'r') as f:
+        config = json.load(f)
+    model_config = OrderedDict()
+    model_config.update(config['pretrained_config'])
+    model_config.update(config['build_config'])
+    return model_config
+
+
+def remove_tensor_padding(input_tensor,
+                          input_tensor_lengths=None,
+                          pad_value=None):
+    if pad_value:
+        assert input_tensor_lengths is None, "input_tensor_lengths should be None when pad_value is provided"
+        # Text tensor case: batch, seq_len
+        assert torch.all(
+            input_tensor[:, 0] != pad_value
+        ), "First token in each sequence should not be pad_value"
+        assert input_tensor_lengths is None
+
+        # Create a mask for all non-pad tokens
+        mask = input_tensor != pad_value
+
+        # Apply the mask to input_tensor to remove pad tokens
+        output_tensor = input_tensor[mask].view(1, -1)
+
+    else:
+        # Audio tensor case: batch, seq_len, feature_len
+        # position_ids case: batch, seq_len
+        assert input_tensor_lengths is not None, "input_tensor_lengths must be provided for 3D input_tensor"
+
+        # Initialize a list to collect valid sequences
+        valid_sequences = []
+
+        for i in range(input_tensor.shape[0]):
+            valid_length = input_tensor_lengths[i]
+            valid_sequences.append(input_tensor[i, :valid_length])
+
+        # Concatenate all valid sequences along the batch dimension
+        output_tensor = torch.cat(valid_sequences, dim=0)
+    return output_tensor
+
+
+class WhisperEncoding:
+
+    def __init__(self, engine_dir):
+        self.session = self.get_session(engine_dir)
+        config = read_config('encoder', engine_dir)
+        self.n_mels = config['n_mels']
+        self.dtype = config['dtype']
+        self.num_languages = config['num_languages']
+        self.encoder_config = config
+
+    def get_session(self, engine_dir):
+        serialize_path = engine_dir / 'encoder' / 'rank0.engine'
+        with open(serialize_path, 'rb') as f:
+            session = Session.from_serialized_engine(f.read())
+        return session
+
+    def get_audio_features(self,
+                           mel,
+                           mel_input_lengths,
+                           encoder_downsampling_factor=2):
+        if isinstance(mel, list):
+            longest_mel = max([f.shape[-1] for f in mel])
+            mel = [
+                torch.nn.functional.pad(f, (0, longest_mel - f.shape[-1]),
+                                        mode='constant') for f in mel
+            ]
+            mel = torch.cat(mel, dim=0).type(
+                str_dtype_to_torch("float16")).contiguous()
+        bsz, seq_len = mel.shape[0], mel.shape[2]
+        position_ids = torch.arange(
+            math.ceil(seq_len / encoder_downsampling_factor),
+            dtype=torch.int32,
+            device=mel.device).expand(bsz, -1).contiguous()
+        if self.encoder_config['plugin_config']['remove_input_padding']:
+            # mel B,D,T -> B,T,D -> BxT, D
+            mel = mel.transpose(1, 2)
+            mel = remove_tensor_padding(mel, mel_input_lengths)
+            position_ids = remove_tensor_padding(
+                position_ids, mel_input_lengths // encoder_downsampling_factor)
+        inputs = OrderedDict()
+        inputs['input_features'] = mel
+        inputs['input_lengths'] = mel_input_lengths
+        inputs['position_ids'] = position_ids
+
+        output_list = [
+            TensorInfo('input_features', str_dtype_to_trt(self.dtype),
+                       mel.shape),
+            TensorInfo('input_lengths', str_dtype_to_trt('int32'),
+                       mel_input_lengths.shape),
+            TensorInfo('position_ids', str_dtype_to_trt('int32'),
+                       inputs['position_ids'].shape)
+        ]
+
+        output_info = (self.session).infer_shapes(output_list)
+
+        logger.debug(f'output info {output_info}')
+        outputs = {
+            t.name: torch.empty(tuple(t.shape),
+                                dtype=trt_dtype_to_torch(t.dtype),
+                                device='cuda')
+            for t in output_info
+        }
+        stream = torch.cuda.current_stream()
+        ok = self.session.run(inputs=inputs,
+                              outputs=outputs,
+                              stream=stream.cuda_stream)
+        assert ok, 'Engine execution failed'
+        stream.synchronize()
+        encoder_output = outputs['encoder_output']
+        encoder_output_lengths = mel_input_lengths // encoder_downsampling_factor
+        return encoder_output, encoder_output_lengths
+
+
+class WhisperDecoding:
+
+    def __init__(self, engine_dir, runtime_mapping, debug_mode=False):
+
+        self.decoder_config = read_config('decoder', engine_dir)
+        self.decoder_generation_session = self.get_session(
+            engine_dir, runtime_mapping, debug_mode)
+
+    def get_session(self, engine_dir, runtime_mapping, debug_mode=False):
+        serialize_path = engine_dir / 'decoder' / 'rank0.engine'
+        with open(serialize_path, "rb") as f:
+            decoder_engine_buffer = f.read()
+
+        decoder_model_config = ModelConfig(
+            max_batch_size=self.decoder_config['max_batch_size'],
+            max_beam_width=self.decoder_config['max_beam_width'],
+            num_heads=self.decoder_config['num_attention_heads'],
+            num_kv_heads=self.decoder_config['num_attention_heads'],
+            hidden_size=self.decoder_config['hidden_size'],
+            vocab_size=self.decoder_config['vocab_size'],
+            cross_attention=True,
+            num_layers=self.decoder_config['num_hidden_layers'],
+            gpt_attention_plugin=self.decoder_config['plugin_config']
+            ['gpt_attention_plugin'],
+            remove_input_padding=self.decoder_config['plugin_config']
+            ['remove_input_padding'],
+            kv_cache_type=KVCacheType.PAGED
+            if self.decoder_config['plugin_config']['paged_kv_cache'] == True
+            else KVCacheType.CONTINUOUS,
+            has_position_embedding=self.
+            decoder_config['has_position_embedding'],
+            dtype=self.decoder_config['dtype'],
+            has_token_type_embedding=False,
+        )
+        decoder_generation_session = tensorrt_llm.runtime.GenerationSession(
+            decoder_model_config,
+            decoder_engine_buffer,
+            runtime_mapping,
+            debug_mode=debug_mode)
+
+        return decoder_generation_session
+
+    def generate(self,
+                 decoder_input_ids,
+                 encoder_outputs,
+                 encoder_max_input_length,
+                 encoder_input_lengths,
+                 eot_id,
+                 max_new_tokens=40,
+                 num_beams=1):
+        batch_size = decoder_input_ids.shape[0]
+        decoder_input_lengths = torch.tensor([
+            decoder_input_ids.shape[-1]
+            for _ in range(decoder_input_ids.shape[0])
+        ],
+                                             dtype=torch.int32,
+                                             device='cuda')
+        decoder_max_input_length = torch.max(decoder_input_lengths).item()
+
+        cross_attention_mask = torch.ones([
+            batch_size, decoder_max_input_length + max_new_tokens,
+            encoder_max_input_length
+        ]).int().cuda()
+        # generation config
+        sampling_config = SamplingConfig(end_id=eot_id,
+                                         pad_id=eot_id,
+                                         num_beams=num_beams)
+        self.decoder_generation_session.setup(
+            decoder_input_lengths.size(0),
+            decoder_max_input_length,
+            max_new_tokens,
+            beam_width=num_beams,
+            encoder_max_input_length=encoder_max_input_length)
+
+        torch.cuda.synchronize()
+
+        decoder_input_ids = decoder_input_ids.type(torch.int32).cuda()
+        if self.decoder_config['plugin_config']['remove_input_padding']:
+            # 50256 is the index of <pad> for all whisper models' decoder
+            WHISPER_PAD_TOKEN_ID = 50256
+            decoder_input_ids = remove_tensor_padding(
+                decoder_input_ids, pad_value=WHISPER_PAD_TOKEN_ID)
+            if encoder_outputs.dim() == 3:
+                encoder_output_lens = torch.full((encoder_outputs.shape[0], ),
+                                                 encoder_outputs.shape[1],
+                                                 dtype=torch.int32,
+                                                 device='cuda')
+
+                encoder_outputs = remove_tensor_padding(encoder_outputs,
+                                                        encoder_output_lens)
+        output_ids = self.decoder_generation_session.decode(
+            decoder_input_ids,
+            decoder_input_lengths,
+            sampling_config,
+            encoder_output=encoder_outputs,
+            encoder_input_lengths=encoder_input_lengths,
+            cross_attention_mask=cross_attention_mask,
+        )
+        torch.cuda.synchronize()
+
+        # get the list of int from output_ids tensor
+        output_ids = output_ids.cpu().numpy().tolist()
+        return output_ids
+
+
+class WhisperTRTLLM(object):
+
+    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)
+        torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node)
+        engine_dir = Path(engine_dir)
+        encoder_config = read_config('encoder', engine_dir)
+        decoder_config = read_config('decoder', engine_dir)
+        self.n_mels = encoder_config['n_mels']
+        self.num_languages = encoder_config['num_languages']
+        is_multilingual = (decoder_config['vocab_size'] >= 51865)
+
+        self.device = device
+        self.tokenizer = get_tokenizer(
+            is_multilingual,
+            num_languages=self.num_languages,
+            language=language,
+            task=task,
+        )
+
+        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,
+        audio: Union[str, np.ndarray, torch.Tensor],
+        padding: int = 0,
+        return_duration=True
+    ):
+        """
+        Compute the log-Mel spectrogram of
+
+        Parameters
+        ----------
+        audio: Union[str, np.ndarray, torch.Tensor], shape = (*)
+            The path to audio or either a NumPy array or Tensor containing the audio waveform in 16 kHz
+
+        n_mels: int
+            The number of Mel-frequency filters, only 80 and 128 are supported
+
+        padding: int
+            Number of zero samples to pad to the right
+
+        device: Optional[Union[str, torch.device]]
+            If given, the audio tensor is moved to this device before STFT
+
+        Returns
+        -------
+        torch.Tensor, shape = (80 or 128, n_frames)
+            A Tensor that contains the Mel spectrogram
+        """
+        if not torch.is_tensor(audio):
+            if isinstance(audio, str):
+                if audio.endswith('.wav'):
+                    audio, _ = load_audio_wav_format(audio)
+                else:
+                    audio = load_audio(audio)
+            assert isinstance(audio, np.ndarray), f"Unsupported audio type: {type(audio)}"
+            duration = audio.shape[-1] / SAMPLE_RATE
+            audio = pad_or_trim(audio, N_SAMPLES)
+            audio = audio.astype(np.float32)
+            audio = torch.from_numpy(audio)
+
+        if self.device is not None:
+            audio = audio.to(self.device)
+        if padding > 0:
+            audio = F.pad(audio, (0, padding))
+        window = torch.hann_window(N_FFT).to(audio.device)
+        stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)
+        magnitudes = stft[..., :-1].abs()**2
+
+        mel_spec = self.filters @ magnitudes
+
+        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+        log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
+        log_spec = (log_spec + 4.0) / 4.0
+        if return_duration:
+            return log_spec, duration
+        else:
+            return log_spec
+
+    def process_batch(
+            self,
+            mel,
+            mel_input_lengths,
+            text_prefix="<|startoftranscript|><|en|><|transcribe|><|notimestamps|>",
+            num_beams=1,
+            max_new_tokens=96):
+        prompt_id = self.tokenizer.encode(
+            text_prefix, allowed_special=set(self.tokenizer.special_tokens.keys()))
+
+        prompt_id = torch.tensor(prompt_id)
+        batch_size = mel.shape[0]
+        decoder_input_ids = prompt_id.repeat(batch_size, 1)
+        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()
+            texts.append(text)
+        return texts
+
+    def transcribe(
+            self,
+            mel,
+            text_prefix="<|startoftranscript|><|en|><|transcribe|><|notimestamps|>",
+            dtype='float16',
+            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
+        mel = mel.repeat(batch_size, 1, 1)
+        if padding_strategy == "longest":
+            pass
+        else:
+            mel = torch.nn.functional.pad(mel, (0, 3000 - mel.shape[2]))
+        features_input_lengths = torch.full((mel.shape[0], ),
+                                             mel.shape[2],
+                                             dtype=torch.int32,
+                                             device=mel.device)
+
+        predictions = self.process_batch(
+            mel,
+            features_input_lengths,
+            text_prefix,
+            num_beams,
+            max_new_tokens=max_new_tokens
+        )
+        prediction = predictions[0]
+
+        # remove all special tokens in the prediction
+        prediction = re.sub(r'<\|.*?\|>', '', prediction)
+        return prediction.strip()
+
+
+def decode_wav_file(
+        model,
+        mel,
+        text_prefix="<|startoftranscript|><|en|><|transcribe|><|notimestamps|>",
+        dtype='float16',
+        batch_size=1,
+        num_beams=1,
+        normalizer=None,
+        mel_filters_dir=None):
+
+    mel = mel.type(str_dtype_to_torch(dtype))
+    mel = mel.unsqueeze(0)
+    # repeat the mel spectrogram to match the batch size
+    mel = mel.repeat(batch_size, 1, 1)
+    predictions = model.process_batch(mel, text_prefix, num_beams)
+    prediction = predictions[0]
+
+    # remove all special tokens in the prediction
+    prediction = re.sub(r'<\|.*?\|>', '', prediction)
+    if normalizer:
+        prediction = normalizer(prediction)
+
+    return prediction.strip()

whisper_live/transcriber_tensorrt.py

@@ -1,338 +0,0 @@
-import json
-import re
-from collections import OrderedDict
-from pathlib import Path
-from typing import Union
-
-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)
-
-import tensorrt_llm
-import tensorrt_llm.logger as logger
-from tensorrt_llm._utils import (str_dtype_to_torch, str_dtype_to_trt,
-                                 trt_dtype_to_torch)
-from tensorrt_llm.runtime import ModelConfig, SamplingConfig
-from tensorrt_llm.runtime.session import Session, TensorInfo
-
-
-SAMPLE_RATE = 16000
-N_FFT = 400
-HOP_LENGTH = 160
-CHUNK_LENGTH = 30
-N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE  # 480000 samples in a 30-second chunk
-
-
-class WhisperEncoding:
-
-    def __init__(self, engine_dir):
-        self.session = self.get_session(engine_dir)
-
-    def get_session(self, engine_dir):
-        config_path = engine_dir / 'encoder_config.json'
-        with open(config_path, 'r') as f:
-            config = json.load(f)
-
-        use_gpt_attention_plugin = config['plugin_config'][
-            'gpt_attention_plugin']
-        dtype = config['builder_config']['precision']
-        n_mels = config['builder_config']['n_mels']
-        num_languages = config['builder_config']['num_languages']
-
-        self.dtype = dtype
-        self.n_mels = n_mels
-        self.num_languages = num_languages
-
-        serialize_path = engine_dir / f'whisper_encoder_{self.dtype}_tp1_rank0.engine'
-
-        with open(serialize_path, 'rb') as f:
-            session = Session.from_serialized_engine(f.read())
-
-        return session
-
-    def get_audio_features(self, mel):
-
-        input_lengths = torch.tensor(
-            [mel.shape[2] // 2 for _ in range(mel.shape[0])],
-            dtype=torch.int32,
-            device=mel.device)
-
-        inputs = OrderedDict()
-        inputs['x'] = mel
-        inputs['input_lengths'] = input_lengths
-
-        output_list = [
-            TensorInfo('x', str_dtype_to_trt(self.dtype), mel.shape),
-            TensorInfo('input_lengths', str_dtype_to_trt('int32'),
-                       input_lengths.shape)
-        ]
-
-        output_info = (self.session).infer_shapes(output_list)
-
-        logger.debug(f'output info {output_info}')
-        outputs = {
-            t.name: torch.empty(tuple(t.shape),
-                                dtype=trt_dtype_to_torch(t.dtype),
-                                device='cuda')
-            for t in output_info
-        }
-        stream = torch.cuda.current_stream()
-        ok = self.session.run(inputs=inputs,
-                              outputs=outputs,
-                              stream=stream.cuda_stream)
-        assert ok, 'Engine execution failed'
-        stream.synchronize()
-        audio_features = outputs['output']
-        return audio_features
-
-
-class WhisperDecoding:
-
-    def __init__(self, engine_dir, runtime_mapping, debug_mode=False):
-
-        self.decoder_config = self.get_config(engine_dir)
-        self.decoder_generation_session = self.get_session(
-            engine_dir, runtime_mapping, debug_mode)
-
-    def get_config(self, engine_dir):
-        config_path = engine_dir / 'decoder_config.json'
-        with open(config_path, 'r') as f:
-            config = json.load(f)
-        decoder_config = OrderedDict()
-        decoder_config.update(config['plugin_config'])
-        decoder_config.update(config['builder_config'])
-        return decoder_config
-
-    def get_session(self, engine_dir, runtime_mapping, debug_mode=False):
-        dtype = self.decoder_config['precision']
-        serialize_path = engine_dir / f'whisper_decoder_{dtype}_tp1_rank0.engine'
-        with open(serialize_path, "rb") as f:
-            decoder_engine_buffer = f.read()
-
-        decoder_model_config = ModelConfig(
-            max_batch_size=self.decoder_config['max_batch_size'],
-            max_beam_width=self.decoder_config['max_beam_width'],
-            num_heads=self.decoder_config['num_heads'],
-            num_kv_heads=self.decoder_config['num_heads'],
-            hidden_size=self.decoder_config['hidden_size'],
-            vocab_size=self.decoder_config['vocab_size'],
-            num_layers=self.decoder_config['num_layers'],
-            gpt_attention_plugin=self.decoder_config['gpt_attention_plugin'],
-            remove_input_padding=self.decoder_config['remove_input_padding'],
-            cross_attention=self.decoder_config['cross_attention'],
-            has_position_embedding=self.
-            decoder_config['has_position_embedding'],
-            has_token_type_embedding=self.
-            decoder_config['has_token_type_embedding'],
-        )
-        decoder_generation_session = tensorrt_llm.runtime.GenerationSession(
-            decoder_model_config,
-            decoder_engine_buffer,
-            runtime_mapping,
-            debug_mode=debug_mode)
-
-        return decoder_generation_session
-
-    def generate(self,
-                 decoder_input_ids,
-                 encoder_outputs,
-                 eot_id,
-                 max_new_tokens=40,
-                 num_beams=1):
-        encoder_input_lengths = torch.tensor(
-            [encoder_outputs.shape[1] for x in range(encoder_outputs.shape[0])],
-            dtype=torch.int32,
-            device='cuda')
-
-        decoder_input_lengths = torch.tensor([
-            decoder_input_ids.shape[-1]
-            for _ in range(decoder_input_ids.shape[0])
-        ],
-                                             dtype=torch.int32,
-                                             device='cuda')
-        decoder_max_input_length = torch.max(decoder_input_lengths).item()
-
-        cross_attention_mask = torch.ones(
-            [encoder_outputs.shape[0], 1,
-             encoder_outputs.shape[1]]).int().cuda()
-
-        # generation config
-        sampling_config = SamplingConfig(end_id=eot_id,
-                                         pad_id=eot_id,
-                                         num_beams=num_beams)
-        self.decoder_generation_session.setup(
-            decoder_input_lengths.size(0),
-            decoder_max_input_length,
-            max_new_tokens,
-            beam_width=num_beams,
-            encoder_max_input_length=encoder_outputs.shape[1])
-
-        torch.cuda.synchronize()
-
-        decoder_input_ids = decoder_input_ids.type(torch.int32).cuda()
-        output_ids = self.decoder_generation_session.decode(
-            decoder_input_ids,
-            decoder_input_lengths,
-            sampling_config,
-            encoder_output=encoder_outputs,
-            encoder_input_lengths=encoder_input_lengths,
-            cross_attention_mask=cross_attention_mask,
-        )
-        torch.cuda.synchronize()
-
-        # get the list of int from output_ids tensor
-        output_ids = output_ids.cpu().numpy().tolist()
-        return output_ids
-
-
-class WhisperTRTLLM(object):
-
-    def __init__(self, engine_dir, assets_dir=None, device=None, is_multilingual=False,
-                 language="en", task="transcribe"):
-        world_size = 1
-        runtime_rank = tensorrt_llm.mpi_rank()
-        runtime_mapping = tensorrt_llm.Mapping(world_size, runtime_rank)
-        torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node)
-        engine_dir = Path(engine_dir)
-
-        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,
-            num_languages=self.encoder.num_languages,
-            language=language,
-            task=task,
-        )
-        self.filters = mel_filters(self.device, self.encoder.n_mels, assets_dir)
-
-    def log_mel_spectrogram(
-        self,
-        audio: Union[str, np.ndarray, torch.Tensor],
-        padding: int = 0,
-        return_duration=True
-    ):
-        """
-        Compute the log-Mel spectrogram of
-
-        Parameters
-        ----------
-        audio: Union[str, np.ndarray, torch.Tensor], shape = (*)
-            The path to audio or either a NumPy array or Tensor containing the audio waveform in 16 kHz
-
-        n_mels: int
-            The number of Mel-frequency filters, only 80 and 128 are supported
-
-        padding: int
-            Number of zero samples to pad to the right
-
-        device: Optional[Union[str, torch.device]]
-            If given, the audio tensor is moved to this device before STFT
-
-        Returns
-        -------
-        torch.Tensor, shape = (80 or 128, n_frames)
-            A Tensor that contains the Mel spectrogram
-        """
-        if not torch.is_tensor(audio):
-            if isinstance(audio, str):
-                if audio.endswith('.wav'):
-                    audio, _ = load_audio_wav_format(audio)
-                else:
-                    audio = load_audio(audio)
-            assert isinstance(audio, np.ndarray), f"Unsupported audio type: {type(audio)}"
-            duration = audio.shape[-1] / SAMPLE_RATE
-            audio = pad_or_trim(audio, N_SAMPLES)
-            audio = audio.astype(np.float32)
-            audio = torch.from_numpy(audio)
-
-        if self.device is not None:
-            audio = audio.to(self.device)
-        if padding > 0:
-            audio = F.pad(audio, (0, padding))
-        window = torch.hann_window(N_FFT).to(audio.device)
-        stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)
-        magnitudes = stft[..., :-1].abs()**2
-
-        mel_spec = self.filters @ magnitudes
-
-        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
-        log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
-        log_spec = (log_spec + 4.0) / 4.0
-        if return_duration:
-            return log_spec, duration
-        else:
-            return log_spec
-
-    def process_batch(
-            self,
-            mel,
-            text_prefix="<|startoftranscript|><|en|><|transcribe|><|notimestamps|>",
-            num_beams=1):
-        prompt_id = self.tokenizer.encode(
-            text_prefix, allowed_special=set(self.tokenizer.special_tokens.keys()))
-
-        prompt_id = torch.tensor(prompt_id)
-        batch_size = mel.shape[0]
-        decoder_input_ids = prompt_id.repeat(batch_size, 1)
-
-        encoder_output = self.encoder.get_audio_features(mel)
-        output_ids = self.decoder.generate(decoder_input_ids,
-                                           encoder_output,
-                                           self.tokenizer.eot,
-                                           max_new_tokens=96,
-                                           num_beams=num_beams)
-        texts = []
-        for i in range(len(output_ids)):
-            text = self.tokenizer.decode(output_ids[i][0]).strip()
-            texts.append(text)
-        return texts
-
-    def transcribe(
-            self,
-            mel,
-            text_prefix="<|startoftranscript|><|en|><|transcribe|><|notimestamps|>",
-            dtype='float16',
-            batch_size=1,
-            num_beams=1,
-            ):
-        mel = mel.type(str_dtype_to_torch(dtype))
-        mel = mel.unsqueeze(0)
-        predictions = self.process_batch(mel, text_prefix, num_beams)
-        prediction = predictions[0]
-
-        # remove all special tokens in the prediction
-        prediction = re.sub(r'<\|.*?\|>', '', prediction)
-        return prediction.strip()
-
-
-def decode_wav_file(
-        model,
-        mel,
-        text_prefix="<|startoftranscript|><|en|><|transcribe|><|notimestamps|>",
-        dtype='float16',
-        batch_size=1,
-        num_beams=1,
-        normalizer=None,
-        mel_filters_dir=None):
-
-    mel = mel.type(str_dtype_to_torch(dtype))
-    mel = mel.unsqueeze(0)
-    # repeat the mel spectrogram to match the batch size
-    mel = mel.repeat(batch_size, 1, 1)
-    predictions = model.process_batch(mel, text_prefix, num_beams)
-    prediction = predictions[0]
-
-    # remove all special tokens in the prediction
-    prediction = re.sub(r'<\|.*?\|>', '', prediction)
-    if normalizer:
-        prediction = normalizer(prediction)
-
-    return prediction.strip()

whisper_live/utils.py

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

whisper_live/vad.py

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