Bump version v0.6.0

Add lock to thread shared variables updates/reads.

.github/workflows/ci.yml

@@ -1,4 +1,4 @@
-name: CI
+name: Test & Build CI/CD
 
 on:
   push:
@@ -7,46 +7,192 @@ on:
     tags:
       - v*
   pull_request:
-    branches:
-      - main
+    branches: [ main ]
+    types: [opened, synchronize, reopened]
 
 jobs:
-  build-and-push-package:
-    runs-on: ubuntu-latest
+  run-tests:
+    runs-on: ubuntu-22.04
+    strategy:
+      matrix:
+        python-version: [3.8, 3.9, '3.10', 3.11]
+    steps:
+      - uses: actions/checkout@v2
+      
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Cache Python dependencies
+        uses: actions/cache@v2
+        with:
+          path: |
+            ~/.cache/pip
+            !~/.cache/pip/log
+          key: ${{ runner.os }}-pip-${{ matrix.python-version }}-${{ hashFiles('requirements/server.txt', 'requirements/client.txt') }}
+          restore-keys: |
+            ${{ runner.os }}-pip-${{ matrix.python-version }}-
+      
+      - name: Install system dependencies
+        run: sudo apt-get update && sudo apt-get install -y ffmpeg portaudio19-dev
+      
+      - name: Install Python dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements/server.txt --extra-index-url https://download.pytorch.org/whl/cpu
+          pip install -r requirements/client.txt
+      
+      - name: Run tests
+        run: |
+          echo "Running tests with Python ${{ matrix.python-version }}"
+          python -m unittest discover -s tests
+  
+  check-code-format:
+    runs-on: ubuntu-22.04
+    strategy:
+      matrix:
+        python-version: [3.8, 3.9, '3.10', 3.11]
 
     steps:
-    - name: Check Out Repository
-      uses: actions/checkout@v2
+      - uses: actions/checkout@v2
 
-    - name: Set up Python
-      uses: actions/setup-python@v2
-      with:
-        python-version: 3.8
-    
-    - name: Set up FFmpeg
-      uses: FedericoCarboni/setup-ffmpeg@v2
-    
-    - name: Install Additional requirements
-      run: |
-        sudo apt-get -y install portaudio19-dev wget
-      shell: bash
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
 
-    - name: Install Client Requirements
-      run: pip install -r requirements/client.txt
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install flake8
 
-    - name: Install Server Requirements
-      run: pip install -r requirements/server.txt
+      - name: Lint with flake8
+        run: |
+          # stop the build if there are Python syntax errors or undefined names
+          flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+          # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+          flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
 
-    - name: Install Wheel for build 
-      run: pip install wheel twine
-    
-    - name: Build wheel
-      run: |
-        python setup.py sdist bdist_wheel
-    
-    - name: Push package on Test PyPI
-      if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
-      uses: pypa/gh-action-pypi-publish@release/v1
-      with:
-        user: __token__
-        password: ${{ secrets.PYPI_API_TOKEN }}
+  build-and-push-docker-cpu:
+    needs: [run-tests, check-code-format]
+    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: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+
+      - name: Build and push Docker image
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          file: docker/Dockerfile.cpu
+          push: true
+          tags: ghcr.io/collabora/whisperlive-cpu:latest
+  
+  build-and-push-docker-tensorrt:
+    needs: [run-tests, check-code-format]
+    timeout-minutes: 60
+    runs-on: ubuntu-22.04
+    if: github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/'))
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v1
+        with:
+          registry: ghcr.io
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.GHCR_TOKEN }}
+
+      - name: Docker Prune
+        run: docker system prune -af
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+
+      - name: Build and push Docker GPU image
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          file: docker/Dockerfile.tensorrt
+          push: true
+          tags: ghcr.io/collabora/whisperlive-tensorrt:latest
+
+  build-and-push-docker-gpu:
+    needs: [run-tests, check-code-format, build-and-push-docker-cpu]
+    timeout-minutes: 20
+    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.gpu
+          push: true
+          tags: ghcr.io/collabora/whisperlive-gpu:latest
+
+  publish-to-pypi:
+    needs: [run-tests, check-code-format]
+    runs-on: ubuntu-22.04
+    if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Set up Python 3.8
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.8
+
+      - name: Cache Python dependencies
+        uses: actions/cache@v2
+        with:
+          path: |
+            ~/.cache/pip
+            !~/.cache/pip/log
+          key: ubuntu-latest-pip-3.8-${{ hashFiles('requirements/server.txt', 'requirements/client.txt') }}
+          restore-keys: |
+            ubuntu-latest-pip-3.8-
+
+      - name: Install system dependencies
+        run: sudo apt-get update && sudo apt-get install -y ffmpeg portaudio19-dev
+
+      - name: Install Python dependencies
+        run: |
+          pip install -r requirements/server.txt
+          pip install -r requirements/client.txt
+          pip install wheel
+
+      - name: Build package
+        run: python setup.py sdist bdist_wheel
+
+      - name: Publish package to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          user: __token__
+          password: ${{ secrets.PYPI_API_TOKEN }}

Audio-Transcription-Chrome/README.md

@@ -26,9 +26,9 @@ To capture the audio in the current tab, we used the chrome `tabCapture` API to
 ### Options
 When using the Audio Transcription extension, you have the following options:
  - **Use Collabora Server**: We provide a demo server which runs the whisper small model.
- - **Use Multilingual Model**: Enable this option to utilize the multilingual capabilities of OpenAI-whisper.
  - **Language**: Select the target language for transcription or translation. You can choose from a variety of languages supported by OpenAI-whisper.
  - **Task:** Choose the specific task to perform on the audio. You can select either "transcribe" for transcription or "translate" to translate the audio to English.
+ - **Model Size**: Select the whisper model size to run the server with.
 
 ### Getting Started
 - Make sure the transcription server is running properly. To know more about how to start the server, see the [documentation here](https://github.com/collabora/whisper-live).

Audio-Transcription-Chrome/background.js

@@ -156,7 +156,9 @@ async function startCapture(options) {
           port: options.port, 
           multilingual: options.useMultilingual,
           language: options.language,
-          task: options.task
+          task: options.task,
+          modelSize: options.modelSize,
+          useVad: options.useVad,
         },
       });
     } else {
@@ -207,13 +209,3 @@ chrome.runtime.onMessage.addListener(async (message) => {
 });
 
 
-/**
- * Listens for if the tab is reloaded.
- * @param {Object} message - The message received from the runtime.
- */
-chrome.tabs.onUpdated.addListener(async (tabId, changeInfo, tab) => {
-  if (changeInfo.status === 'complete') {
-    await executeScriptInTab(tabId, "content.js");
-    await delayExecution(500);
-  }
-});

Audio-Transcription-Chrome/content.js

@@ -59,7 +59,7 @@ function init_element() {
 
     elem_container = document.createElement('div');
     elem_container.id = "transcription";
-    elem_container.style.cssText = 'padding-top:16px;font-size:18px;line-height:18px;top:0px;position:absolute;width:500px;height:90px;opacity:0.9;z-index:100;background:black;border-radius:10px;color:white;';
+    elem_container.style.cssText = 'padding-top:16px;font-size:18px;position: fixed; top: 85%; left: 50%; transform: translate(-50%, -50%);line-height:18px;width:500px;height:90px;opacity:0.9;z-index:100;background:black;border-radius:10px;color:white;';
 
     for (var i = 0; i < 4; i++) {
         elem_text = document.createElement('span');
@@ -173,13 +173,13 @@ chrome.runtime.onMessage.addListener((request, sender, sendResponse) => {
     if (type === "STOP") {
         remove_element();
         sendResponse({data: "STOPPED"});
-        return;
+        return true;
     } else if (type === "showWaitPopup"){
         initPopupElement();
 
         showPopup(`Estimated wait time ~ ${Math.round(data)} minutes`);
         sendResponse({data: "popup"});
-        return;
+        return true;
     }
 
     init_element();
@@ -234,4 +234,5 @@ chrome.runtime.onMessage.addListener((request, sender, sendResponse) => {
     }
 
     sendResponse({});
+    return true;
 });

Audio-Transcription-Chrome/options.js

@@ -93,16 +93,14 @@ async function startRecord(option) {
     const socket = new WebSocket(`ws://${option.host}:${option.port}/`);
     let isServerReady = false;
     let language = option.language;
-    if (language === null && !option.multilingual) {
-      language = 'en';
-    }
     socket.onopen = function(e) { 
       socket.send(
         JSON.stringify({
           uid: uuid,
-          multilingual: option.multilingual,
           language: option.language,
-          task: option.task
+          task: option.task,
+          model: option.modelSize,
+          use_vad: option.useVad
         })
       );
     };
@@ -184,16 +182,17 @@ async function startRecord(option) {
  * @param {Object} sender - The sender object containing information about the message sender.
  * @param {Function} sendResponse - The function to send a response back to the message sender.
  */
-chrome.runtime.onMessage.addListener(async (request, sender, sendResponse) => {
+chrome.runtime.onMessage.addListener((request, sender, sendResponse) => {
   const { type, data } = request;
 
   switch (type) {
     case "start_capture":
-      await startRecord(data);
+      startRecord(data);
       break;
     default:
       break;
   }
 
   sendResponse({});
+  return true;
 });

Audio-Transcription-Chrome/popup.html

@@ -16,120 +16,137 @@
     <label for="useServerCheckbox">Use Collabora Whisper-Live Server</label>
   </div>
   <div class="checkbox-container">
-    <input type="checkbox" id="useMultilingualCheckbox">
-    <label for="useMultilingualCheckbox">Use Multilingual Model</label>
+    <input type="checkbox" id="useVadCheckbox">
+    <label for="useVadCheckbox">Use Voice Activity Detection</label>
   </div>
   <div class="dropdown-container">
     <label for="languageDropdown">Select Language:</label>
-    <select id="languageDropdown" disabled>
-      <option value="">Select Language</option>
-      <option value="zh">Chinese</option>
-      <option value="de">German</option>
-      <option value="es">Spanish</option>
-      <option value="ru">Russian</option>
-      <option value="ko">Korean</option>
-      <option value="fr">French</option>
-      <option value="ja">Japanese</option>
-      <option value="pt">Portuguese</option>
-      <option value="tr">Turkish</option>
-      <option value="pl">Polish</option>
-      <option value="ca">Catalan</option>
-      <option value="nl">Dutch</option>
-      <option value="ar">Arabic</option>
-      <option value="sv">Swedish</option>
-      <option value="it">Italian</option>
-      <option value="id">Indonesian</option>
-      <option value="hi">Hindi</option>
-      <option value="fi">Finnish</option>
-      <option value="vi">Vietnamese</option>
-      <option value="he">Hebrew</option>
-      <option value="uk">Ukrainian</option>
-      <option value="el">Greek</option>
-      <option value="ms">Malay</option>
-      <option value="cs">Czech</option>
-      <option value="ro">Romanian</option>
-      <option value="da">Danish</option>
-      <option value="hu">Hungarian</option>
-      <option value="ta">Tamil</option>
-      <option value="no">Norwegian</option>
-      <option value="th">Thai</option>
-      <option value="ur">Urdu</option>
-      <option value="hr">Croatian</option>
-      <option value="bg">Bulgarian</option>
-      <option value="lt">Lithuanian</option>
-      <option value="la">Latin</option>
-      <option value="mi">Maori</option>
-      <option value="ml">Malayalam</option>
-      <option value="cy">Welsh</option>
-      <option value="sk">Slovak</option>
-      <option value="te">Telugu</option>
-      <option value="fa">Persian</option>
-      <option value="lv">Latvian</option>
-      <option value="bn">Bengali</option>
-      <option value="sr">Serbian</option>
-      <option value="az">Azerbaijani</option>
-      <option value="sl">Slovenian</option>
-      <option value="kn">Kannada</option>
-      <option value="et">Estonian</option>
-      <option value="mk">Macedonian</option>
-      <option value="br">Breton</option>
-      <option value="eu">Basque</option>
-      <option value="is">Icelandic</option>
-      <option value="hy">Armenian</option>
-      <option value="ne">Nepali</option>
-      <option value="mn">Mongolian</option>
-      <option value="bs">Bosnian</option>
-      <option value="kk">Kazakh</option>
-      <option value="sq">Albanian</option>
-      <option value="sw">Swahili</option>
-      <option value="gl">Galician</option>
-      <option value="mr">Marathi</option>
-      <option value="pa">Punjabi</option>
-      <option value="si">Sinhala</option>
-      <option value="km">Khmer</option>
-      <option value="sn">Shona</option>
-      <option value="yo">Yoruba</option>
-      <option value="so">Somali</option>
+    <select id="languageDropdown">
+      <option value="" selected>Automatically detect</option>
       <option value="af">Afrikaans</option>
-      <option value="oc">Occitan</option>
-      <option value="ka">Georgian</option>
-      <option value="be">Belarusian</option>
-      <option value="tg">Tajik</option>
-      <option value="sd">Sindhi</option>
-      <option value="gu">Gujarati</option>
+      <option value="sq">Albanian</option>
       <option value="am">Amharic</option>
-      <option value="yi">Yiddish</option>
-      <option value="lo">Lao</option>
-      <option value="uz">Uzbek</option>
-      <option value="fo">Faroese</option>
-      <option value="ht">Haitian Creole</option>
-      <option value="ps">Pashto</option>
-      <option value="tk">Turkmen</option>
-      <option value="nn">Nynorsk</option>
-      <option value="mt">Maltese</option>
-      <option value="sa">Sanskrit</option>
-      <option value="lb">Luxembourgish</option>
-      <option value="my">Myanmar</option>
-      <option value="bo">Tibetan</option>
-      <option value="tl">Tagalog</option>
-      <option value="mg">Malagasy</option>
+      <option value="ar">Arabic</option>
+      <option value="hy">Armenian</option>
       <option value="as">Assamese</option>
-      <option value="tt">Tatar</option>
-      <option value="haw">Hawaiian</option>
-      <option value="ln">Lingala</option>
-      <option value="ha">Hausa</option>
+      <option value="az">Azerbaijani</option>
       <option value="ba">Bashkir</option>
+      <option value="eu">Basque</option>
+      <option value="be">Belarusian</option>
+      <option value="bn">Bengali</option>
+      <option value="bs">Bosnian</option>
+      <option value="br">Breton</option>
+      <option value="bg">Bulgarian</option>
+      <option value="ca">Catalan</option>
+      <option value="zh">Chinese</option>
+      <option value="hr">Croatian</option>
+      <option value="cs">Czech</option>
+      <option value="da">Danish</option>
+      <option value="nl">Dutch</option>
+      <option value="en">English</option>
+      <option value="et">Estonian</option>
+      <option value="fo">Faroese</option>
+      <option value="fi">Finnish</option>
+      <option value="fr">French</option>
+      <option value="gl">Galician</option>
+      <option value="ka">Georgian</option>
+      <option value="de">German</option>
+      <option value="el">Greek</option>
+      <option value="gu">Gujarati</option>
+      <option value="ht">Haitian Creole</option>
+      <option value="ha">Hausa</option>
+      <option value="haw">Hawaiian</option>
+      <option value="he">Hebrew</option>
+      <option value="hi">Hindi</option>
+      <option value="hu">Hungarian</option>
+      <option value="is">Icelandic</option>
+      <option value="id">Indonesian</option>
+      <option value="it">Italian</option>
+      <option value="ja">Japanese</option>
       <option value="jw">Javanese</option>
+      <option value="kn">Kannada</option>
+      <option value="kk">Kazakh</option>
+      <option value="km">Khmer</option>
+      <option value="ko">Korean</option>
+      <option value="lo">Lao</option>
+      <option value="la">Latin</option>
+      <option value="lv">Latvian</option>
+      <option value="ln">Lingala</option>
+      <option value="lt">Lithuanian</option>
+      <option value="lb">Luxembourgish</option>
+      <option value="mk">Macedonian</option>
+      <option value="mg">Malagasy</option>
+      <option value="ms">Malay</option>
+      <option value="ml">Malayalam</option>
+      <option value="mt">Maltese</option>
+      <option value="mi">Maori</option>
+      <option value="mr">Marathi</option>
+      <option value="mn">Mongolian</option>
+      <option value="my">Myanmar</option>
+      <option value="ne">Nepali</option>
+      <option value="no">Norwegian</option>
+      <option value="nn">Nynorsk</option>
+      <option value="oc">Occitan</option>
+      <option value="ps">Pashto</option>
+      <option value="fa">Persian</option>
+      <option value="pl">Polish</option>
+      <option value="pt">Portuguese</option>
+      <option value="pa">Punjabi</option>
+      <option value="ro">Romanian</option>
+      <option value="ru">Russian</option>
+      <option value="sa">Sanskrit</option>
+      <option value="sr">Serbian</option>
+      <option value="sn">Shona</option>
+      <option value="sd">Sindhi</option>
+      <option value="si">Sinhala</option>
+      <option value="sk">Slovak</option>
+      <option value="sl">Slovenian</option>
+      <option value="so">Somali</option>
+      <option value="es">Spanish</option>
       <option value="su">Sundanese</option>
+      <option value="sw">Swahili</option>
+      <option value="sv">Swedish</option>
+      <option value="tl">Tagalog</option>
+      <option value="tg">Tajik</option>
+      <option value="ta">Tamil</option>
+      <option value="tt">Tatar</option>
+      <option value="te">Telugu</option>
+      <option value="th">Thai</option>
+      <option value="bo">Tibetan</option>
+      <option value="tr">Turkish</option>
+      <option value="tk">Turkmen</option>
+      <option value="uk">Ukrainian</option>
+      <option value="ur">Urdu</option>
+      <option value="uz">Uzbek</option>
+      <option value="vi">Vietnamese</option>
+      <option value="cy">Welsh</option>
+      <option value="yi">Yiddish</option>
+      <option value="yo">Yoruba</option>
     </select>
   </div>
   <div class="dropdown-container">
     <label for="taskDropdown">Select task:</label>
-    <select id="taskDropdown" disabled>
+    <select id="taskDropdown" >
       <option value="">Select Task</option>
       <option value="transcribe" selected>Transcribe</option>
       <option value="translate">Translate</option>
     </select>
   </div>
+  <div class="dropdown-container">
+    <label for="modelSizeDropdown">Select Model Size:</label>
+    <select id="modelSizeDropdown">
+      <option value="">Select model</option>
+      <option value="tiny">Tiny </option>
+      <option value="tiny.en">Tiny (English-only)</option>
+      <option value="base">Base</option>
+      <option value="base.en">Base (English-only)</option>
+      <option value="small" selected>Small</option>
+      <option value="small.en">Small (English-only)</option>
+      <option value="medium">Medium</option>
+      <option value="medium.en">Medium (English-only)</option>
+      <option value="large-v2">Large-v2</option>
+      <option value="large-v3">Large-v3</option>
+    </select>
+  </div>
 </body>
 </html>

Audio-Transcription-Chrome/popup.js

@@ -4,11 +4,13 @@ document.addEventListener("DOMContentLoaded", function () {
   const stopButton = document.getElementById("stopCapture");
 
   const useServerCheckbox = document.getElementById("useServerCheckbox");
-  const useMultilingualCheckbox = document.getElementById('useMultilingualCheckbox');
+  const useVadCheckbox = document.getElementById("useVadCheckbox");
   const languageDropdown = document.getElementById('languageDropdown');
   const taskDropdown = document.getElementById('taskDropdown');
+  const modelSizeDropdown = document.getElementById('modelSizeDropdown');
   let selectedLanguage = null;
   let selectedTask = taskDropdown.value;
+  let selectedModelSize = modelSizeDropdown.value;
 
   // Add click event listeners to the buttons
   startButton.addEventListener("click", startCapture);
@@ -30,11 +32,9 @@ document.addEventListener("DOMContentLoaded", function () {
     }
   });
 
-  chrome.storage.local.get("useMultilingualModelState", ({ useMultilingualModelState }) => {
-    if (useMultilingualModelState !== undefined) {
-      useMultilingualCheckbox.checked = useMultilingualModelState;
-      languageDropdown.disabled = !useMultilingualModelState;
-      taskDropdown.disabled = !useMultilingualModelState;
+  chrome.storage.local.get("useVadState", ({ useVadState }) => {
+    if (useVadState !== undefined) {
+      useVadCheckbox.checked = useVadState;
     }
   });
 
@@ -52,6 +52,13 @@ document.addEventListener("DOMContentLoaded", function () {
     }
   });
 
+  chrome.storage.local.get("selectedModelSize", ({ selectedModelSize: storedModelSize }) => {
+    if (storedModelSize !== undefined) {
+      modelSizeDropdown.value = storedModelSize;
+      selectedModelSize = storedModelSize;
+    }
+  });
+
   // Function to handle the start capture button click event
   async function startCapture() {
     // Ignore click if the button is disabled
@@ -77,9 +84,10 @@ document.addEventListener("DOMContentLoaded", function () {
         tabId: currentTab.id,
         host: host,
         port: port,
-        useMultilingual: useMultilingualCheckbox.checked,
         language: selectedLanguage,
-        task: selectedTask
+        task: selectedTask,
+        modelSize: selectedModelSize,
+        useVad: useVadCheckbox.checked,
       }, () => {
         // Update capturing state in storage and toggle the buttons
         chrome.storage.local.set({ capturingState: { isCapturing: true } }, () => {
@@ -118,9 +126,11 @@ document.addEventListener("DOMContentLoaded", function () {
   function toggleCaptureButtons(isCapturing) {
     startButton.disabled = isCapturing;
     stopButton.disabled = !isCapturing;
-    useServerCheckbox.disabled = isCapturing; 
-    useMultilingualCheckbox.disabled = isCapturing;
-    
+    useServerCheckbox.disabled = isCapturing;
+    useVadCheckbox.disabled = isCapturing;
+    modelSizeDropdown.disabled = isCapturing;
+    languageDropdown.disabled = isCapturing;
+    taskDropdown.disabled = isCapturing; 
     startButton.classList.toggle("disabled", isCapturing);
     stopButton.classList.toggle("disabled", !isCapturing);
   }
@@ -131,16 +141,9 @@ document.addEventListener("DOMContentLoaded", function () {
     chrome.storage.local.set({ useServerState });
   });
 
-  useMultilingualCheckbox.addEventListener('change', function() {
-    const useMultilingualModelState = useMultilingualCheckbox.checked;
-    if (useMultilingualModelState) {
-      languageDropdown.disabled = false;
-      taskDropdown.disabled = false;
-    } else {
-      languageDropdown.disabled = true;
-      taskDropdown.disabled = true;
-    }
-    chrome.storage.local.set({ useMultilingualModelState });
+  useVadCheckbox.addEventListener("change", () => {
+    const useVadState = useVadCheckbox.checked;
+    chrome.storage.local.set({ useVadState });
   });
 
   languageDropdown.addEventListener('change', function() {
@@ -157,6 +160,11 @@ document.addEventListener("DOMContentLoaded", function () {
     chrome.storage.local.set({ selectedTask });
   });
 
+  modelSizeDropdown.addEventListener('change', function() {
+    selectedModelSize = modelSizeDropdown.value;
+    chrome.storage.local.set({ selectedModelSize });
+  });
+
   chrome.runtime.onMessage.addListener(async (request, sender, sendResponse) => {
     if (request.action === "updateSelectedLanguage") {
       const detectedLanguage = request.detectedLanguage;

Audio-Transcription-Firefox/README.md

@@ -24,9 +24,9 @@ To capture the audio in the current tab, we used the chrome `tabCapture` API to
 ### Options
 When using the Audio Transcription extension, you have the following options:
  - **Use Collabora Server**: We provide a demo server which runs the whisper small model.
- - **Use Multilingual Model**: Enable this option to utilize the multilingual capabilities of OpenAI-whisper.
  - **Language**: Select the target language for transcription or translation. You can choose from a variety of languages supported by OpenAI-whisper.
  - **Task:** Choose the specific task to perform on the audio. You can select either "transcribe" for transcription or "translate" to translate the audio to English.
+  - **Model Size**: Select the whisper model size to run the server with.
 
 ### Getting Started
 - Make sure the transcription server is running properly. To know more about how to start the server, see the [documentation here](https://github.com/collabora/whisper-live).

Audio-Transcription-Firefox/content.js

@@ -66,18 +66,16 @@ function resampleTo16kHZ(audioData, origSampleRate = 44100) {
 function startRecording(data) {
     socket = new WebSocket(`ws://${data.host}:${data.port}/`);
     language = data.language;
-    if (language === null && !data.useMultilingual) {
-      language = 'en';
-    }
 
     const uuid = generateUUID();
     socket.onopen = function(e) { 
       socket.send(
         JSON.stringify({
             uid: uuid,
-            multilingual: data.useMultilingual,
             language: data.language,
-            task: data.task
+            task: data.task,
+            model: data.modelSize,
+            use_vad: data.useVad
         })
       );
     };
@@ -200,7 +198,7 @@ function init_element() {
 
     elem_container = document.createElement('div');
     elem_container.id = "transcription";
-    elem_container.style.cssText = 'padding-top:16px;font-size:18px;line-height:18px;top:0px;position:absolute;width:500px;height:90px;opacity:0.9;z-index:100;background:black;border-radius:10px;color:white;';
+    elem_container.style.cssText = 'padding-top:16px;font-size:18px;line-height:18px;position:fixed;top:85%;left:50%;transform:translate(-50%,-50%);width:500px;height:90px;opacity:0.9;z-index:100;background:black;border-radius:10px;color:white;';
 
     for (var i = 0; i < 4; i++) {
         elem_text = document.createElement('span');

Audio-Transcription-Firefox/popup.html

@@ -15,114 +15,114 @@
     <input type="checkbox" id="useServerCheckbox">
     <label for="useServerCheckbox">Use Collabora Whisper-Live Server</label>
   </div>
-  <textarea id="waitTextBox" style="display: none;"></textarea>
-
   <div class="checkbox-container">
-    <input type="checkbox" id="useMultilingualCheckbox">
-    <label for="useMultilingualCheckbox">Use Multilingual Model</label>
+    <input type="checkbox" id="useVadCheckbox">
+    <label for="useVadCheckbox">Use Voice Activity Detection</label>
   </div>
+  <textarea id="waitTextBox" style="display: none;"></textarea>
   <div class="dropdown-container">
     <label for="languageDropdown">Select Language:</label>
-    <select id="languageDropdown" disabled>
-      <option value="">Select Language</option>
-      <option value="zh">Chinese</option>
-      <option value="de">German</option>
-      <option value="es">Spanish</option>
-      <option value="ru">Russian</option>
-      <option value="ko">Korean</option>
-      <option value="fr">French</option>
-      <option value="ja">Japanese</option>
-      <option value="pt">Portuguese</option>
-      <option value="tr">Turkish</option>
-      <option value="pl">Polish</option>
-      <option value="ca">Catalan</option>
-      <option value="nl">Dutch</option>
-      <option value="ar">Arabic</option>
-      <option value="sv">Swedish</option>
-      <option value="it">Italian</option>
-      <option value="id">Indonesian</option>
-      <option value="hi">Hindi</option>
-      <option value="fi">Finnish</option>
-      <option value="vi">Vietnamese</option>
-      <option value="he">Hebrew</option>
-      <option value="uk">Ukrainian</option>
-      <option value="el">Greek</option>
-      <option value="ms">Malay</option>
-      <option value="cs">Czech</option>
-      <option value="ro">Romanian</option>
-      <option value="da">Danish</option>
-      <option value="hu">Hungarian</option>
-      <option value="ta">Tamil</option>
-      <option value="no">Norwegian</option>
-      <option value="th">Thai</option>
-      <option value="ur">Urdu</option>
-      <option value="hr">Croatian</option>
-      <option value="bg">Bulgarian</option>
-      <option value="lt">Lithuanian</option>
-      <option value="la">Latin</option>
-      <option value="mi">Maori</option>
-      <option value="ml">Malayalam</option>
-      <option value="cy">Welsh</option>
-      <option value="sk">Slovak</option>
-      <option value="te">Telugu</option>
-      <option value="fa">Persian</option>
-      <option value="lv">Latvian</option>
-      <option value="bn">Bengali</option>
-      <option value="sr">Serbian</option>
-      <option value="az">Azerbaijani</option>
-      <option value="sl">Slovenian</option>
-      <option value="kn">Kannada</option>
-      <option value="et">Estonian</option>
-      <option value="mk">Macedonian</option>
-      <option value="br">Breton</option>
-      <option value="eu">Basque</option>
-      <option value="is">Icelandic</option>
-      <option value="hy">Armenian</option>
-      <option value="ne">Nepali</option>
-      <option value="mn">Mongolian</option>
-      <option value="bs">Bosnian</option>
-      <option value="kk">Kazakh</option>
-      <option value="sq">Albanian</option>
-      <option value="sw">Swahili</option>
-      <option value="gl">Galician</option>
-      <option value="mr">Marathi</option>
-      <option value="pa">Punjabi</option>
-      <option value="si">Sinhala</option>
-      <option value="km">Khmer</option>
-      <option value="sn">Shona</option>
-      <option value="yo">Yoruba</option>
-      <option value="so">Somali</option>
+    <select id="languageDropdown">
+      <option value="" selected>Automatically detect</option>
       <option value="af">Afrikaans</option>
-      <option value="oc">Occitan</option>
-      <option value="ka">Georgian</option>
-      <option value="be">Belarusian</option>
-      <option value="tg">Tajik</option>
-      <option value="sd">Sindhi</option>
-      <option value="gu">Gujarati</option>
+      <option value="sq">Albanian</option>
       <option value="am">Amharic</option>
-      <option value="yi">Yiddish</option>
-      <option value="lo">Lao</option>
-      <option value="uz">Uzbek</option>
-      <option value="fo">Faroese</option>
-      <option value="ht">Haitian Creole</option>
-      <option value="ps">Pashto</option>
-      <option value="tk">Turkmen</option>
-      <option value="nn">Nynorsk</option>
-      <option value="mt">Maltese</option>
-      <option value="sa">Sanskrit</option>
-      <option value="lb">Luxembourgish</option>
-      <option value="my">Myanmar</option>
-      <option value="bo">Tibetan</option>
-      <option value="tl">Tagalog</option>
-      <option value="mg">Malagasy</option>
+      <option value="ar">Arabic</option>
+      <option value="hy">Armenian</option>
       <option value="as">Assamese</option>
-      <option value="tt">Tatar</option>
-      <option value="haw">Hawaiian</option>
-      <option value="ln">Lingala</option>
-      <option value="ha">Hausa</option>
+      <option value="az">Azerbaijani</option>
       <option value="ba">Bashkir</option>
+      <option value="eu">Basque</option>
+      <option value="be">Belarusian</option>
+      <option value="bn">Bengali</option>
+      <option value="bs">Bosnian</option>
+      <option value="br">Breton</option>
+      <option value="bg">Bulgarian</option>
+      <option value="ca">Catalan</option>
+      <option value="zh">Chinese</option>
+      <option value="hr">Croatian</option>
+      <option value="cs">Czech</option>
+      <option value="da">Danish</option>
+      <option value="nl">Dutch</option>
+      <option value="en">English</option>
+      <option value="et">Estonian</option>
+      <option value="fo">Faroese</option>
+      <option value="fi">Finnish</option>
+      <option value="fr">French</option>
+      <option value="gl">Galician</option>
+      <option value="ka">Georgian</option>
+      <option value="de">German</option>
+      <option value="el">Greek</option>
+      <option value="gu">Gujarati</option>
+      <option value="ht">Haitian Creole</option>
+      <option value="ha">Hausa</option>
+      <option value="haw">Hawaiian</option>
+      <option value="he">Hebrew</option>
+      <option value="hi">Hindi</option>
+      <option value="hu">Hungarian</option>
+      <option value="is">Icelandic</option>
+      <option value="id">Indonesian</option>
+      <option value="it">Italian</option>
+      <option value="ja">Japanese</option>
       <option value="jw">Javanese</option>
+      <option value="kn">Kannada</option>
+      <option value="kk">Kazakh</option>
+      <option value="km">Khmer</option>
+      <option value="ko">Korean</option>
+      <option value="lo">Lao</option>
+      <option value="la">Latin</option>
+      <option value="lv">Latvian</option>
+      <option value="ln">Lingala</option>
+      <option value="lt">Lithuanian</option>
+      <option value="lb">Luxembourgish</option>
+      <option value="mk">Macedonian</option>
+      <option value="mg">Malagasy</option>
+      <option value="ms">Malay</option>
+      <option value="ml">Malayalam</option>
+      <option value="mt">Maltese</option>
+      <option value="mi">Maori</option>
+      <option value="mr">Marathi</option>
+      <option value="mn">Mongolian</option>
+      <option value="my">Myanmar</option>
+      <option value="ne">Nepali</option>
+      <option value="no">Norwegian</option>
+      <option value="nn">Nynorsk</option>
+      <option value="oc">Occitan</option>
+      <option value="ps">Pashto</option>
+      <option value="fa">Persian</option>
+      <option value="pl">Polish</option>
+      <option value="pt">Portuguese</option>
+      <option value="pa">Punjabi</option>
+      <option value="ro">Romanian</option>
+      <option value="ru">Russian</option>
+      <option value="sa">Sanskrit</option>
+      <option value="sr">Serbian</option>
+      <option value="sn">Shona</option>
+      <option value="sd">Sindhi</option>
+      <option value="si">Sinhala</option>
+      <option value="sk">Slovak</option>
+      <option value="sl">Slovenian</option>
+      <option value="so">Somali</option>
+      <option value="es">Spanish</option>
       <option value="su">Sundanese</option>
+      <option value="sw">Swahili</option>
+      <option value="sv">Swedish</option>
+      <option value="tl">Tagalog</option>
+      <option value="tg">Tajik</option>
+      <option value="ta">Tamil</option>
+      <option value="tt">Tatar</option>
+      <option value="te">Telugu</option>
+      <option value="th">Thai</option>
+      <option value="bo">Tibetan</option>
+      <option value="tr">Turkish</option>
+      <option value="tk">Turkmen</option>
+      <option value="uk">Ukrainian</option>
+      <option value="ur">Urdu</option>
+      <option value="uz">Uzbek</option>
+      <option value="vi">Vietnamese</option>
+      <option value="cy">Welsh</option>
+      <option value="yi">Yiddish</option>
+      <option value="yo">Yoruba</option>
     </select>
   </div>
   <div class="dropdown-container">
@@ -133,5 +133,21 @@
       <option value="translate">Translate</option>
     </select>
   </div>
+  <div class="dropdown-container">
+    <label for="modelSizeDropdown">Select Model Size:</label>
+    <select id="modelSizeDropdown">
+      <option value="">Select model</option>
+      <option value="tiny">Tiny </option>
+      <option value="tiny.en">Tiny (English-only)</option>
+      <option value="base">Base</option>
+      <option value="base.en">Base (English-only)</option>
+      <option value="small" selected>Small</option>
+      <option value="small.en">Small (English-only)</option>
+      <option value="medium">Medium</option>
+      <option value="medium.en">Medium (English-only)</option>
+      <option value="large-v2">Large-v2</option>
+      <option value="large-v3">Large-v3</option>
+    </select>
+  </div>
 </body>
-</html>
+</html>

Audio-Transcription-Firefox/popup.js

@@ -3,11 +3,14 @@ document.addEventListener("DOMContentLoaded", function() {
   const stopButton = document.getElementById("stopCapture");
 
   const useServerCheckbox = document.getElementById("useServerCheckbox");
-  const useMultilingualCheckbox = document.getElementById('useMultilingualCheckbox');
+  const useVadCheckbox = document.getElementById("useVadCheckbox");
   const languageDropdown = document.getElementById('languageDropdown');
   const taskDropdown = document.getElementById('taskDropdown');
+  const modelSizeDropdown = document.getElementById('modelSizeDropdown');
   let selectedLanguage = null;
   let selectedTask = taskDropdown.value;
+  let selectedModelSize = modelSizeDropdown.value;
+  
 
   browser.storage.local.get("capturingState")
     .then(function(result) {
@@ -32,11 +35,9 @@ document.addEventListener("DOMContentLoaded", function() {
     }
   });
 
-  browser.storage.local.get("useMultilingualModelState", ({ useMultilingualModelState }) => {
-    if (useMultilingualModelState !== undefined) {
-      useMultilingualCheckbox.checked = useMultilingualModelState;
-      languageDropdown.disabled = !useMultilingualModelState;
-      taskDropdown.disabled = !useMultilingualModelState;
+  browser.storage.local.get("useVadState", ({ useVadState }) => {
+    if (useVadState !== undefined) {
+      useVadCheckbox.checked = useVadState;
     }
   });
 
@@ -54,6 +55,13 @@ document.addEventListener("DOMContentLoaded", function() {
     }
   });
 
+  browser.storage.local.get("selectedModelSize", ({ selectedModelSize: storedModelSize }) => {
+    if (storedModelSize !== undefined) {
+      modelSizeDropdown.value = storedModelSize;
+      selectedModelSize = storedModelSize;
+    }
+  });
+
   startButton.addEventListener("click", function() {
     let host = "localhost";
     let port = "9090";
@@ -73,9 +81,10 @@ document.addEventListener("DOMContentLoaded", function() {
             data: {
               host: host,
               port: port,
-              useMultilingual: useMultilingualCheckbox.checked,
               language: selectedLanguage,
-              task: selectedTask
+              task: selectedTask,
+              modelSize: selectedModelSize,
+              useVad: useVadCheckbox.checked,
             } 
           });
         toggleCaptureButtons(true);
@@ -114,8 +123,10 @@ document.addEventListener("DOMContentLoaded", function() {
     startButton.disabled = isCapturing;
     stopButton.disabled = !isCapturing;
     useServerCheckbox.disabled = isCapturing;
-    useMultilingualCheckbox.disabled = isCapturing;
-
+    useVadCheckbox.disabled = isCapturing;
+    modelSizeDropdown.disabled = isCapturing;
+    languageDropdown.disabled = isCapturing;
+    taskDropdown.disabled = isCapturing; 
     startButton.classList.toggle("disabled", isCapturing);
     stopButton.classList.toggle("disabled", !isCapturing);
   }
@@ -126,16 +137,9 @@ document.addEventListener("DOMContentLoaded", function() {
     browser.storage.local.set({ useServerState });
   });
 
-  useMultilingualCheckbox.addEventListener('change', function() {
-    const useMultilingualModelState = useMultilingualCheckbox.checked;
-    if (useMultilingualModelState) {
-      languageDropdown.disabled = false;
-      taskDropdown.disabled = false;
-    } else {
-      languageDropdown.disabled = true;
-      taskDropdown.disabled = true;
-    }
-    browser.storage.local.set({ useMultilingualModelState });
+  useVadCheckbox.addEventListener("change", () => {
+    const useVadState = useVadCheckbox.checked;
+    browser.storage.local.set({ useVadState });
   });
 
   languageDropdown.addEventListener('change', function() {
@@ -152,6 +156,11 @@ document.addEventListener("DOMContentLoaded", function() {
     browser.storage.local.set({ selectedTask });
   });
 
+  modelSizeDropdown.addEventListener('change', function() {
+    selectedModelSize = modelSizeDropdown.value;
+    browser.storage.local.set({ selectedModelSize });
+  });
+
   browser.runtime.onMessage.addListener((request, sender, sendResponse) => {
     if (request.action === "updateSelectedLanguage") {
       const detectedLanguage = request.data;

Audio-Transcription-Firefox/style.css

@@ -108,4 +108,4 @@ label {
 
 .dropdown-container {
   padding: 10px; 
-}
+}

README.md

@@ -1,14 +1,20 @@
-# whisper-live
-A nearly-live implementation of OpenAI's Whisper.
+# WhisperLive
 
-This project is a real-time transcription application that uses the OpenAI Whisper model to convert speech input into text output. It can be used to transcribe both live audio input from microphone and pre-recorded audio files.
+<h2 align="center">
+  <a href="https://www.youtube.com/watch?v=0PHWCApIcCI"><img
+src="https://img.youtube.com/vi/0PHWCApIcCI/0.jpg" style="background-color:rgba(0,0,0,0);" height=300 alt="WhisperLive"></a>
+  <br><br>A nearly-live implementation of OpenAI's Whisper.
+<br><br>
+</h2>
 
-Unlike traditional speech recognition systems that rely on continuous audio streaming, we use [voice activity detection (VAD)](https://github.com/snakers4/silero-vad) to detect the presence of speech and only send the audio data to whisper when speech is detected. This helps to reduce the amount of data sent to the whisper model and improves the accuracy of the transcription output.
+This project is a real-time transcription application that uses the OpenAI Whisper model
+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
 - Install PyAudio and ffmpeg
 ```bash
- bash setup.sh
+ bash scripts/setup.sh
 ```
 
 - Install whisper-live from pip
@@ -16,64 +22,138 @@ Unlike traditional speech recognition systems that rely on continuous audio stre
  pip install whisper-live
 ```
 
+### Setting up NVIDIA/TensorRT-LLM for TensorRT backend
+- 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
-- Run the server
-```python
- from whisper_live.server import TranscriptionServer
- server = TranscriptionServer()
- server.run("0.0.0.0", 9090)
+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)
+
+### Running the Server
+- [Faster Whisper](https://github.com/SYSTRAN/faster-whisper) backend
+```bash
+python3 run_server.py --port 9090 \
+                      --backend faster_whisper
+  
+# running with custom model
+python3 run_server.py --port 9090 \
+                      --backend faster_whisper \
+                      -fw "/path/to/custom/faster/whisper/model"
 ```
 
-- On the client side
-    - To transcribe an audio file:
-    ```python
-      from whisper_live.client import TranscriptionClient
-      client = TranscriptionClient("localhost", 9090, is_multilingual=True, lang="hi", translate=True)
-      client(audio_file_path)
-    ```
-    This command transcribes the specified audio file (audio.wav) using the Whisper model. It connects to the server running on localhost at port 9090. It also enables the multilingual feature, allowing transcription in multiple languages. The language option specifies the target language for transcription, in this case, Hindi ("hi"). 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.
+- TensorRT backend. Currently, we recommend to only use the docker setup for TensorRT. Follow [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md) which works as expected. Make sure to build your TensorRT Engines before running the server with TensorRT backend.
+```bash
+# Run English only model
+python3 run_server.py -p 9090 \
+                      -b tensorrt \
+                      -trt /home/TensorRT-LLM/examples/whisper/whisper_small_en
 
-    - To transcribe from microphone:
-    ```python
-      from whisper_live.client import TranscriptionClient
-      client = TranscriptionClient(host, port, is_multilingual=True, lang="hi", translate=True)
-      client()
-    ```
-    This command captures audio from the microphone and sends it to the server for transcription. It uses the same options as the previous command, enabling the multilingual feature and specifying the target language and task.
-
-
-## Transcribe audio from browser
-- Run the server
-```python
- from whisper_live.server import TranscriptionServer
- server = TranscriptionServer()
- server.run("0.0.0.0", 9090)
+# Run Multilingual model
+python3 run_server.py -p 9090 \
+                      -b tensorrt \
+                      -trt /home/TensorRT-LLM/examples/whisper/whisper_small \
+                      -m
+```
+#### 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
+python3 run_server.py --port 9090 \
+                      --backend faster_whisper \
+                      --omp_num_threads 4
 ```
-This would start the websocket server on port ```9090```.
 
-### Chrome Extension
-- Refer to [Audio-Transcription-Chrome](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Chrome#readme) to use Chrome extension.
+#### Single model mode
+By default, when running the server without specifying a model, the server will instantiate a new whisper model for every client connection. This has the advantage, that the server can use different model sizes, based on the client's requested model size. On the other hand, it also means you have to wait for the model to be loaded upon client connection and you will have increased (V)RAM usage.
 
-### Firefox Extension
-- Refer to [Audio-Transcription-Firefox](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Firefox#readme) to use Mozilla Firefox extension.
+When serving a custom TensorRT model using the `-trt` or a custom faster_whisper model using the `-fw` option, the server will instead only instantiate the custom model once and then reuse it for all client connections.
+
+If you don't want this, set `--no_single_model`.
+
+
+### Running the Client
+- Initializing the client with below parameters:
+  - `lang`: Language of the input audio, applicable only if using a multilingual model.
+  - `translate`: If set to `True` then translate from any language to `en`.
+  - `model`: Whisper model size.
+  - `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.
+
+```python
+from whisper_live.client import TranscriptionClient
+client = TranscriptionClient(
+  "localhost",
+  9090,
+  lang="en",
+  translate=False,
+  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
+  max_clients=4,
+  max_connection_time=600
+)
+```
+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.
+
+- Transcribe an audio file:
+```python
+client("tests/jfk.wav")
+```
+
+- To transcribe from microphone:
+```python
+client()
+```
+
+- To transcribe from a RTSP stream:
+```python
+client(rtsp_url="rtsp://admin:admin@192.168.0.1/rtsp")
+```
+
+- To transcribe from a HLS stream:
+```python
+client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/bbc_1xtra.isml/bbc_1xtra-audio%3d96000.norewind.m3u8")
+```
+
+## 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.
 
 ## Whisper Live Server in Docker
 - GPU
-```bash
- docker build . -t whisper-live -f docker/Dockerfile.gpu
- docker run -it --gpus all -p 9090:9090 whisper-live:latest
-```
+  - Faster-Whisper
+  ```bash
+  docker run -it --gpus all -p 9090:9090 ghcr.io/collabora/whisperlive-gpu:latest
+  ```
+
+  - TensorRT. 
+  ```bash
+  docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt
+
+  # Build small.en engine
+  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en        # float16
+  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int8   # int8 weight only quantization
+  bash build_whisper_tensorrt.sh /app/TensorRT-LLM-examples small.en int4   # int4 weight only quantization
+
+  # Run server with small.en
+  python3 run_server.py --port 9090 \
+                        --backend tensorrt \
+                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_float16"
+                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int8"
+                        --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_en_int4"
+  ```
 
 - CPU
 ```bash
- docker build . -t whisper-live -f docker/Dockerfile.cpu
- docker run -it -p 9090:9090 whisper-live:latest
+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.
-- [ ] TensorRT backend for Whisper.
+- [x] TensorRT backend for Whisper.
 
 ## Contact
 
@@ -98,6 +178,5 @@ We are available to help you with both Open Source and proprietary AI projects.
   publisher = {GitHub},
   journal = {GitHub repository},
   howpublished = {\url{https://github.com/snakers4/silero-vad}},
-  commit = {insert_some_commit_here},
   email = {hello@silero.ai}
 }

TensorRT_whisper.md

@@ -0,0 +1,38 @@
+# WhisperLive-TensorRT
+We have only tested the TensorRT backend in docker so, we recommend docker for a smooth TensorRT backend setup.
+**Note**: We use `tensorrt_llm==0.15.0.dev2024111200`
+
+## 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)
+
+- 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
+```
+
+## 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        # 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
+```
+
+## Run WhisperLive Server with TensorRT Backend
+```bash
+# Run English only model
+python3 run_server.py --port 9090 \
+                      --backend tensorrt \
+                      --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_float16" \
+                      --trt_multilingual
+```

docker/Dockerfile.cpu

@@ -1,45 +1,25 @@
-FROM ubuntu:focal
+FROM python:3.10-bookworm
 
 ARG DEBIAN_FRONTEND=noninteractive
 
-# Remove any third-party apt sources to avoid issues with expiring keys.
-RUN rm -f /etc/apt/sources.list.d/*.list
+# install lib required for pyaudio
+RUN apt update && apt install -y portaudio19-dev && apt-get clean && rm -rf /var/lib/apt/lists/*
 
-# Install some basic utilities.
-RUN apt-get update && apt-get install -y \
-    curl \
-    ca-certificates \
-    sudo \
-    git \
-    bzip2 \
-    libx11-6 \
- && rm -rf /var/lib/apt/lists/*
+# update pip to support for whl.metadata -> less downloading
+RUN pip install --no-cache-dir -U "pip>=24"
 
-RUN apt update
-
-# install python
-RUN apt install software-properties-common -y && \
-    add-apt-repository ppa:deadsnakes/ppa && \
-    apt update
-
-RUN apt install python3-dev -y && \
-    apt install python-is-python3
-    
-
-# install pip
-RUN apt install python3-pip -y
-
-# Create a working directory.
+# create a working directory
 RUN mkdir /app
 WORKDIR /app
 
-COPY setup.sh /app
-COPY requirements/ /app
+# install pytorch, but without the nvidia-libs that are only necessary for gpu
+RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu
 
-RUN bash setup.sh
-RUN pip install -r server.txt
+# install the requirements for running the whisper-live server
+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"]

docker/Dockerfile.gpu

@@ -1,47 +1,26 @@
-FROM nvidia/cuda:11.2.2-cudnn8-runtime-ubuntu20.04
+FROM python:3.10-bookworm
 
 ARG DEBIAN_FRONTEND=noninteractive
 
-# Remove any third-party apt sources to avoid issues with expiring keys.
-RUN rm -f /etc/apt/sources.list.d/*.list
+# install lib required for pyaudio
+RUN apt update && apt install -y portaudio19-dev && apt-get clean && rm -rf /var/lib/apt/lists/*
 
-# Install some basic utilities.
-RUN apt-get update && apt-get install -y \
-    curl \
-    ca-certificates \
-    sudo \
-    git \
-    bzip2 \
-    libx11-6 \
- && rm -rf /var/lib/apt/lists/*
+# update pip to support for whl.metadata -> less downloading
+RUN pip install --no-cache-dir -U "pip>=24"
 
-RUN apt update
-
-# install python
-RUN apt install software-properties-common -y && \
-    add-apt-repository ppa:deadsnakes/ppa && \
-    apt update
-
-RUN apt install python3-dev -y && \
-    apt install python-is-python3
-    
-
-# install pip
-RUN apt install python3-pip -y
-
-# Create a working directory.
+# create a working directory
 RUN mkdir /app
 WORKDIR /app
 
-COPY setup.sh /app
-COPY requirements/ /app
+# install the requirements for running the whisper-live server
+COPY requirements/server.txt /app/
+RUN pip install --no-cache-dir -r server.txt && rm server.txt
 
-RUN apt update --fix-missing
-RUN bash setup.sh
-RUN pip install -r server.txt
+# make the paths of the nvidia libs installed as wheels visible. equivalent to:
+# export LD_LIBRARY_PATH=`python3 -c 'import os; import nvidia.cublas.lib; import nvidia.cudnn.lib; print(os.path.dirname(nvidia.cublas.lib.__file__) + ":" + os.path.dirname(nvidia.cudnn.lib.__file__))'`
+ENV LD_LIBRARY_PATH="/usr/local/lib/python3.10/site-packages/nvidia/cublas/lib:/usr/local/lib/python3.10/site-packages/nvidia/cudnn/lib"
 
 COPY whisper_live /app/whisper_live
-
 COPY run_server.py /app
 
 CMD ["python", "run_server.py"]

docker/Dockerfile.tensorrt

@@ -0,0 +1,30 @@
+FROM nvidia/cuda:12.5.1-runtime-ubuntu22.04 AS base
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update && apt-get install -y \
+    python3.10 python3-pip openmpi-bin libopenmpi-dev git git-lfs wget \
+    && rm -rf /var/lib/apt/lists/*
+
+FROM base AS devel
+RUN pip3 install --no-cache-dir -U tensorrt_llm==0.15.0.dev2024111200 --extra-index-url https://pypi.nvidia.com
+WORKDIR /app
+RUN git clone https://github.com/NVIDIA/TensorRT-LLM.git && cd TensorRT-LLM && \
+    git checkout c629546ce429623c8a163633095230154a6f0574 && cd ../ && \
+    mv TensorRT-LLM/examples ./TensorRT-LLM-examples && \
+    rm -rf TensorRT-LLM
+
+
+FROM devel AS release
+WORKDIR /app
+COPY assets/ ./assets
+RUN wget -nc -P assets/ https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz
+
+COPY scripts/setup.sh ./
+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/server.txt

@@ -1,7 +1,13 @@
-PyAudio
-faster-whisper==0.6.0
---extra-index-url https://download.pytorch.org/whl/cu111
-torch==1.10.1
-torchaudio==0.10.1
+faster-whisper==1.1.0
 websockets
-onnxruntime==1.16.0
+onnxruntime==1.16.0
+numba
+kaldialign
+soundfile
+ffmpeg-python
+scipy
+jiwer
+evaluate
+numpy<2
+openai-whisper==20240930
+tokenizers==0.20.3

run_server.py

@@ -1,5 +1,50 @@
-from whisper_live.server import TranscriptionServer
+import argparse
+import os
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--port', '-p',
+                        type=int,
+                        default=9090,
+                        help="Websocket port to run the server on.")
+    parser.add_argument('--backend', '-b',
+                        type=str,
+                        default='faster_whisper',
+                        help='Backends from ["tensorrt", "faster_whisper"]')
+    parser.add_argument('--faster_whisper_custom_model_path', '-fw',
+                        type=str, default=None,
+                        help="Custom Faster Whisper Model")
+    parser.add_argument('--trt_model_path', '-trt',
+                        type=str,
+                        default=None,
+                        help='Whisper TensorRT model path')
+    parser.add_argument('--trt_multilingual', '-m',
+                        action="store_true",
+                        help='Boolean only for TensorRT model. True if multilingual.')
+    parser.add_argument('--omp_num_threads', '-omp',
+                        type=int,
+                        default=1,
+                        help="Number of threads to use for OpenMP")
+    parser.add_argument('--no_single_model', '-nsm',
+                        action='store_true',
+                        help='Set this if every connection should instantiate its own model. Only relevant for custom model, passed using -trt or -fw.')
+    args = parser.parse_args()
+
+    if args.backend == "tensorrt":
+        if args.trt_model_path is None:
+            raise ValueError("Please Provide a valid tensorrt model path")
+
+    if "OMP_NUM_THREADS" not in os.environ:
+        os.environ["OMP_NUM_THREADS"] = str(args.omp_num_threads)
+
+    from whisper_live.server import TranscriptionServer
     server = TranscriptionServer()
-    server.run("0.0.0.0")
+    server.run(
+        "0.0.0.0",
+        port=args.port,
+        backend=args.backend,
+        faster_whisper_custom_model_path=args.faster_whisper_custom_model_path,
+        whisper_tensorrt_path=args.trt_model_path,
+        trt_multilingual=args.trt_multilingual,
+        single_model=not args.no_single_model,
+    )

scripts/build_whisper_tensorrt.sh

@@ -0,0 +1,122 @@
+#!/bin/bash
+
+download_and_build_model() {
+    local model_name="$1"
+    local model_url=""
+
+    case "$model_name" in
+        "tiny.en")
+            model_url="https://openaipublic.azureedge.net/main/whisper/models/d3dd57d32accea0b295c96e26691aa14d8822fac7d9d27d5dc00b4ca2826dd03/tiny.en.pt"
+            ;;
+        "tiny")
+            model_url="https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt"
+            ;;
+        "base.en")
+            model_url="https://openaipublic.azureedge.net/main/whisper/models/25a8566e1d0c1e2231d1c762132cd20e0f96a85d16145c3a00adf5d1ac670ead/base.en.pt"
+            ;;
+        "base")
+            model_url="https://openaipublic.azureedge.net/main/whisper/models/ed3a0b6b1c0edf879ad9b11b1af5a0e6ab5db9205f891f668f8b0e6c6326e34e/base.pt"
+            ;;
+        "small.en")
+            model_url="https://openaipublic.azureedge.net/main/whisper/models/f953ad0fd29cacd07d5a9eda5624af0f6bcf2258be67c92b79389873d91e0872/small.en.pt"
+            ;;
+        "small")
+            model_url="https://openaipublic.azureedge.net/main/whisper/models/9ecf779972d90ba49c06d968637d720dd632c55bbf19d441fb42bf17a411e794/small.pt"
+            ;;
+        "medium.en")
+            model_url="https://openaipublic.azureedge.net/main/whisper/models/d7440d1dc186f76616474e0ff0b3b6b879abc9d1a4926b7adfa41db2d497ab4f/medium.en.pt"
+            ;;
+        "medium")
+            model_url="https://openaipublic.azureedge.net/main/whisper/models/345ae4da62f9b3d59415adc60127b97c714f32e89e936602e85993674d08dcb1/medium.pt"
+            ;;
+        "large-v1")
+            model_url="https://openaipublic.azureedge.net/main/whisper/models/e4b87e7e0bf463eb8e6956e646f1e277e901512310def2c24bf0e11bd3c28e9a/large-v1.pt"
+            ;;
+        "large-v2")
+            model_url="https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt"
+            ;;
+        "large-v3" | "large")
+            model_url="https://openaipublic.azureedge.net/main/whisper/models/e5b1a55b89c1367dacf97e3e19bfd829a01529dbfdeefa8caeb59b3f1b81dadb/large-v3.pt"
+            ;;
+        "large-v3-turbo" | "turbo")
+            model_url="https://openaipublic.azureedge.net/main/whisper/models/aff26ae408abcba5fbf8813c21e62b0941638c5f6eebfb145be0c9839262a19a/large-v3-turbo.pt"
+            ;;
+        *)
+            echo "Invalid model name: $model_name"
+            exit 1
+            ;;
+    esac
+
+    if [ "$model_name" == "turbo" ]; then
+        model_name="large-v3-turbo"
+    fi
+
+    local inference_precision="float16"
+    local weight_only_precision="${2:-float16}"
+    local max_beam_width=4
+    local max_batch_size=1
+
+    echo "Downloading $model_name..."
+    # wget --directory-prefix=assets "$model_url"
+    # echo "Download completed: ${model_name}.pt"
+    if [ ! -f "assets/${model_name}.pt" ]; then
+        wget --directory-prefix=assets "$model_url"
+        echo "Download completed: ${model_name}.pt"
+    else
+        echo "${model_name}.pt already exists in assets directory."
+    fi
+
+    local sanitized_model_name="${model_name//./_}"
+    local checkpoint_dir="whisper_${sanitized_model_name}_weights_${weight_only_precision}"
+    local output_dir="whisper_${sanitized_model_name}_${weight_only_precision}"
+    echo "$output_dir"
+    echo "Converting model weights for $model_name..."
+    python3 convert_checkpoint.py \
+        $( [[ "$weight_only_precision" == "int8" || "$weight_only_precision" == "int4" ]] && echo "--use_weight_only --weight_only_precision $weight_only_precision" ) \
+        --output_dir "$checkpoint_dir" --model_name "$model_name"
+    
+    echo "Building encoder for $model_name..."
+    trtllm-build \
+        --checkpoint_dir "${checkpoint_dir}/encoder" \
+        --output_dir "${output_dir}/encoder" \
+        --moe_plugin disable \
+        --enable_xqa disable \
+        --max_batch_size "$max_batch_size" \
+        --gemm_plugin disable \
+        --bert_attention_plugin "$inference_precision" \
+        --max_input_len 3000 \
+        --max_seq_len 3000
+    
+    echo "Building decoder for $model_name..."
+    trtllm-build \
+        --checkpoint_dir "${checkpoint_dir}/decoder" \
+        --output_dir "${output_dir}/decoder" \
+        --moe_plugin disable \
+        --enable_xqa disable \
+        --max_beam_width "$max_beam_width" \
+        --max_batch_size "$max_batch_size" \
+        --max_seq_len 200 \
+        --max_input_len 14 \
+        --max_encoder_input_len 3000 \
+        --gemm_plugin "$inference_precision" \
+        --bert_attention_plugin "$inference_precision" \
+        --gpt_attention_plugin "$inference_precision"
+
+    echo "TensorRT LLM engine built for $model_name."
+    echo "========================================="
+    echo "Model is located at: $(pwd)/$output_dir"
+}
+
+if [ "$#" -lt 1 ]; then
+    echo "Usage: $0 <path-to-tensorrt-examples-dir> [model-name]"
+    exit 1
+fi
+
+tensorrt_examples_dir="$1"
+model_name="${2:-small.en}"
+weight_only_precision="${3:-float16}"  # Default to float16 if not provided
+
+cd $tensorrt_examples_dir/whisper
+pip install --no-deps -r requirements.txt
+
+download_and_build_model "$model_name" "$weight_only_precision"

setup.py

@@ -10,45 +10,52 @@ HERE = pathlib.Path(__file__).parent
 README = (HERE / "README.md").read_text()
 
 # This call to setup() does all the work
-setup(name="whisper-live",
-      version=__version__,
-      description="A nearly-live implementation of OpenAI's Whisper.",
-      long_description=README,
-      long_description_content_type="text/markdown",
-      include_package_data=True,
-      url="https://github.com/collabora/WhisperLive",
-      author="Collabora Ltd",
-      author_email="vineet.suryan@collabora.com",
-      license="MIT",
-      classifiers=[
-          "Development Status :: 4 - Beta",
-          "Intended Audience :: Developers",
-          "Intended Audience :: Science/Research",
-          "License :: OSI Approved :: MIT License",
-          "Programming Language :: Python :: 3",
-          "Programming Language :: Python :: 3 :: Only",
-          "Programming Language :: Python :: 3.8",
-          "Programming Language :: Python :: 3.9",
-          "Topic :: Scientific/Engineering :: Artificial Intelligence",
-      ],
-      packages=find_packages(
-        exclude=("examples",
-                 "Audio-Transcription-Chrome",
-                 "Audio-Transcription-Firefox",
-                 "requirements",
-                 "whisper-finetuning"
-                )
-      ),
-      install_requires=[
+setup(
+    name="whisper-live",
+    version=__version__,
+    description="A nearly-live implementation of OpenAI's Whisper.",
+    long_description=README,
+    long_description_content_type="text/markdown",
+    include_package_data=True,
+    url="https://github.com/collabora/WhisperLive",
+    author="Collabora Ltd",
+    author_email="vineet.suryan@collabora.com",
+    license="MIT",
+    classifiers=[
+        "Development Status :: 4 - Beta",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Science/Research",
+        "License :: OSI Approved :: MIT License",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3 :: Only",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    ],
+    packages=find_packages(
+        exclude=(
+            "examples",
+            "Audio-Transcription-Chrome",
+            "Audio-Transcription-Firefox",
+            "requirements",
+            "whisper-finetuning"
+        )
+    ),
+    install_requires=[
         "PyAudio",
-        "faster-whisper==0.6.0",
+        "faster-whisper==1.1.0",
         "torch",
         "torchaudio",
         "websockets",
-        "onnxruntime",
+        "onnxruntime==1.16.0",
         "ffmpeg-python",
         "scipy",
         "websocket-client",
-      ],
-      python_requires=">=3.8"
-)
+        "numba",
+        "openai-whisper==20240930",
+        "kaldialign",
+        "soundfile",
+        "tokenizers==0.20.3"
+    ],
+    python_requires=">=3.8"
+)

tests/test_client.py

@@ -0,0 +1,158 @@
+import json
+import os
+import scipy
+import websocket
+import copy
+import unittest
+from unittest.mock import patch, MagicMock
+from whisper_live.client import Client, TranscriptionClient, TranscriptionTeeClient
+from whisper_live.utils import resample
+from pathlib import Path
+
+
+class BaseTestCase(unittest.TestCase):
+    @patch('whisper_live.client.websocket.WebSocketApp')
+    @patch('whisper_live.client.pyaudio.PyAudio')
+    def setUp(self, mock_pyaudio, mock_websocket):
+        self.mock_pyaudio_instance = MagicMock()
+        mock_pyaudio.return_value = self.mock_pyaudio_instance
+        self.mock_stream = MagicMock()
+        self.mock_pyaudio_instance.open.return_value = self.mock_stream
+
+        self.mock_ws_app = mock_websocket.return_value
+        self.mock_ws_app.send = MagicMock()
+
+        self.client = TranscriptionClient(host='localhost', port=9090, lang="en").client
+
+        self.mock_pyaudio = mock_pyaudio
+        self.mock_websocket = mock_websocket
+        self.mock_audio_packet = b'\x00\x01\x02\x03'
+
+    def tearDown(self):
+        self.client.close_websocket()
+        self.mock_pyaudio.stop()
+        self.mock_websocket.stop()
+        del self.client
+
+class TestClientWebSocketCommunication(BaseTestCase):
+    def test_websocket_communication(self):
+        expected_url = 'ws://localhost:9090'
+        self.mock_websocket.assert_called()
+        self.assertEqual(self.mock_websocket.call_args[0][0], expected_url)
+
+
+class TestClientCallbacks(BaseTestCase):
+    def test_on_open(self):
+        expected_message = json.dumps({
+            "uid": self.client.uid,
+            "language": self.client.language,
+            "task": self.client.task,
+            "model": self.client.model,
+            "use_vad": True,
+            "max_clients": 4,
+            "max_connection_time": 600,
+        })
+        self.client.on_open(self.mock_ws_app)
+        self.mock_ws_app.send.assert_called_with(expected_message)
+
+    def test_on_message(self):
+        message = json.dumps(
+            {
+                "uid": self.client.uid,
+                "message": "SERVER_READY",
+                "backend": "faster_whisper"
+            }
+        )
+        self.client.on_message(self.mock_ws_app, message)
+
+        message = json.dumps({
+            "uid": self.client.uid,
+            "segments": [
+                {"start": 0, "end": 1, "text": "Test transcript", "completed": True},
+                {"start": 1, "end": 2, "text": "Test transcript 2", "completed": True},
+                {"start": 2, "end": 3, "text": "Test transcript 3", "completed": True}
+            ]
+        })
+        self.client.on_message(self.mock_ws_app, message)
+
+        # Assert that the transcript was updated correctly
+        self.assertEqual(len(self.client.transcript), 3)
+        self.assertEqual(self.client.transcript[1]['text'], "Test transcript 2")
+
+    def test_on_close(self):
+        close_status_code = 1000
+        close_msg = "Normal closure"
+        self.client.on_close(self.mock_ws_app, close_status_code, close_msg)
+
+        self.assertFalse(self.client.recording)
+        self.assertFalse(self.client.server_error)
+        self.assertFalse(self.client.waiting)
+
+    def test_on_error(self):
+        error_message = "Test Error"
+        self.client.on_error(self.mock_ws_app, error_message)
+
+        self.assertTrue(self.client.server_error)
+        self.assertEqual(self.client.error_message, error_message)
+
+
+class TestAudioResampling(unittest.TestCase):
+    def test_resample_audio(self):
+        original_audio = "assets/jfk.flac"
+        expected_sr = 16000
+        resampled_audio = resample(original_audio, expected_sr)
+
+        sr, _ = scipy.io.wavfile.read(resampled_audio)
+        self.assertEqual(sr, expected_sr)
+
+        os.remove(resampled_audio)
+
+
+class TestSendingAudioPacket(BaseTestCase):
+    def test_send_packet(self):
+        self.client.send_packet_to_server(self.mock_audio_packet)
+        self.client.client_socket.send.assert_called_with(self.mock_audio_packet, websocket.ABNF.OPCODE_BINARY)
+
+class TestTee(BaseTestCase):
+    @patch('whisper_live.client.websocket.WebSocketApp')
+    @patch('whisper_live.client.pyaudio.PyAudio')
+    def setUp(self, mock_audio, mock_websocket):
+        super().setUp()
+        self.client2 = Client(host='localhost', port=9090, lang="es", translate=False, srt_file_path="transcript.srt")
+        self.client3 = Client(host='localhost', port=9090, lang="es", translate=True, srt_file_path="translation.srt")
+        # need a separate mock for each websocket
+        self.client3.client_socket = copy.deepcopy(self.client3.client_socket)
+        self.tee = TranscriptionTeeClient([self.client2, self.client3])
+
+    def tearDown(self):
+        self.tee.close_all_clients()
+        del self.tee
+        super().tearDown()
+
+    def test_invalid_constructor(self):
+        with self.assertRaises(Exception) as context:
+            TranscriptionTeeClient([])
+
+    def test_multicast_unconditional(self):
+        self.tee.multicast_packet(self.mock_audio_packet, True)
+        for client in self.tee.clients:
+            client.client_socket.send.assert_called_with(self.mock_audio_packet, websocket.ABNF.OPCODE_BINARY)
+
+    def test_multicast_conditional(self):
+        self.client2.recording = False
+        self.client3.recording = True
+        self.tee.multicast_packet(self.mock_audio_packet, False)
+        self.client2.client_socket.send.assert_not_called()
+        self.client3.client_socket.send.assert_called_with(self.mock_audio_packet, websocket.ABNF.OPCODE_BINARY)
+
+    def test_close_all(self):
+        self.tee.close_all_clients()
+        for client in self.tee.clients:
+            client.client_socket.close.assert_called()
+
+    def test_write_all_srt(self):
+        for client in self.tee.clients:
+            client.server_backend = "faster_whisper"
+        self.tee.write_all_clients_srt()
+        self.assertTrue(Path("transcript.srt").is_file())
+        self.assertTrue(Path("translation.srt").is_file())

tests/test_server.py

@@ -0,0 +1,148 @@
+import subprocess
+import time
+import json
+import unittest
+from unittest import mock
+
+import numpy as np
+import jiwer
+
+from websockets.exceptions import ConnectionClosed
+from whisper_live.server import TranscriptionServer, BackendType, ClientManager
+from whisper_live.client import Client, TranscriptionClient, TranscriptionTeeClient
+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, {})
+        self.assertDictEqual(server.client_manager.start_times, {})
+
+
+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
+        }
+        self.server.client_manager.max_connection_time = 600
+
+    def test_get_wait_time(self):
+        expected_wait_time = (600 - (time.time() - self.server.client_manager.start_times['client2'])) / 60
+        print(self.server.client_manager.get_wait_time(), expected_wait_time)
+        self.assertAlmostEqual(self.server.client_manager.get_wait_time(), expected_wait_time, places=2)
+
+
+class TestServerConnection(unittest.TestCase):
+    def setUp(self):
+        self.server = TranscriptionServer()
+
+    @mock.patch('websockets.WebSocketCommonProtocol')
+    def test_connection(self, mock_websocket):
+        mock_websocket.recv.return_value = json.dumps({
+            'uid': 'test_client',
+            'language': 'en',
+            'task': 'transcribe',
+            'model': 'tiny.en'
+        })
+        self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
+
+    @mock.patch('websockets.WebSocketCommonProtocol')
+    def test_recv_audio_exception_handling(self, mock_websocket):
+        mock_websocket.recv.side_effect = [json.dumps({
+            'uid': 'test_client',
+            'language': 'en',
+            'task': 'transcribe',
+            'model': 'tiny.en'
+        }),  np.array([1, 2, 3]).tobytes()]
+
+        with self.assertLogs(level="ERROR"):
+            self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
+
+        self.assertNotIn(mock_websocket, self.server.client_manager.clients)
+
+
+class TestServerInferenceAccuracy(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.mock_pyaudio_patch = mock.patch('pyaudio.PyAudio')
+        cls.mock_pyaudio = cls.mock_pyaudio_patch.start()
+        cls.mock_pyaudio.return_value.open.return_value = mock.MagicMock()
+        
+        cls.server_process = subprocess.Popen(["python", "run_server.py"])
+        time.sleep(2)
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.server_process.terminate()
+        cls.server_process.wait()
+
+    def setUp(self):
+        self.normalizer = EnglishTextNormalizer()
+
+    def check_prediction(self, srt_path):
+        gt = "And so my fellow Americans, ask not, what your country can do for you. Ask what you can do for your country!"
+        with open(srt_path, "r") as f:
+            lines = f.readlines()
+            prediction = " ".join([line.strip() for line in lines[2::4]])
+        prediction_normalized = self.normalizer(prediction)
+        gt_normalized = self.normalizer(gt)
+
+        # calculate WER
+        wer_score = jiwer.wer(gt_normalized, prediction_normalized)
+        self.assertLess(wer_score, 0.05)
+
+    def test_inference(self):
+        client = TranscriptionClient(
+            "localhost", "9090", model="base.en", lang="en",
+        )
+        client("assets/jfk.flac")
+        self.check_prediction("output.srt")
+
+    def test_simultaneous_inference(self):
+        client1 = Client(
+            "localhost", "9090", model="base.en", lang="en", srt_file_path="transcript1.srt")
+        client2 = Client(
+            "localhost", "9090", model="base.en", lang="en", srt_file_path="transcript2.srt")
+        tee = TranscriptionTeeClient([client1, client2])
+        tee("assets/jfk.flac")
+        self.check_prediction("transcript1.srt")
+        self.check_prediction("transcript2.srt")
+
+
+class TestExceptionHandling(unittest.TestCase):
+    def setUp(self):
+        self.server = TranscriptionServer()
+
+    @mock.patch('websockets.WebSocketCommonProtocol')
+    def test_connection_closed_exception(self, mock_websocket):
+        mock_websocket.recv.side_effect = ConnectionClosed(1001, "testing connection closed", rcvd_then_sent=mock.Mock())
+
+        with self.assertLogs(level="INFO") as log:
+            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')
+    def test_json_decode_exception(self, mock_websocket):
+        mock_websocket.recv.return_value = "invalid json"
+
+        with self.assertLogs(level="ERROR") as log:
+            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')
+    def test_unexpected_exception_handling(self, mock_websocket):
+        mock_websocket.recv.side_effect = RuntimeError("Unexpected error")
+
+        with self.assertLogs(level="ERROR") as log:
+            self.server.recv_audio(mock_websocket, BackendType("faster_whisper"))
+            for message in log.output:
+                print(message)
+            print()
+            self.assertTrue(any("Unexpected error" in message for message in log.output))

tests/test_vad.py

@@ -0,0 +1,26 @@
+import unittest
+import numpy as np
+from whisper_live.tensorrt_utils import load_audio
+from whisper_live.vad import VoiceActivityDetector
+
+
+class TestVoiceActivityDetection(unittest.TestCase):
+    def setUp(self):
+        self.vad = VoiceActivityDetector()
+        self.sample_rate = 16000
+
+    def generate_silence(self, duration_seconds):
+        return np.zeros(int(self.sample_rate * duration_seconds), dtype=np.float32)
+
+    def load_speech_segment(self, filepath):
+        return load_audio(filepath)
+
+    def test_vad_silence_detection(self):
+        silence = self.generate_silence(3)
+        is_speech_present = self.vad(silence.copy())
+        self.assertFalse(is_speech_present, "VAD incorrectly identified silence as speech.")
+
+    def test_vad_speech_detection(self):
+        audio_tensor = load_audio("assets/jfk.flac")
+        is_speech_present = self.vad(audio_tensor)
+        self.assertTrue(is_speech_present, "VAD failed to identify speech segment.")

whisper_live/__version__.py

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

whisper_live/client.py

@@ -1,56 +1,38 @@
 import os
+import shutil
 import wave
 
+import logging
 import numpy as np
-import scipy
-import ffmpeg
 import pyaudio
 import threading
-import textwrap
 import json
 import websocket
 import uuid
 import time
-
-
-def resample(file: str, sr: int = 16000):
-    """
-    # https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/audio.py#L22
-    Open an audio file and read as mono waveform, resampling as necessary,
-    save the resampled audio
-
-    Args:
-        file (str): The audio file to open
-        sr (int): The sample rate to resample the audio if necessary
-    
-    Returns:
-        resampled_file (str): The resampled audio file
-    """
-    try:
-        # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
-        # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
-        out, _ = (
-            ffmpeg.input(file, threads=0)
-            .output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=sr)
-            .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
-        )
-    except ffmpeg.Error as e:
-        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
-    np_buffer = np.frombuffer(out, dtype=np.int16)
-
-    resampled_file = f"{file.split('.')[0]}_resampled.wav"
-    scipy.io.wavfile.write(resampled_file, sr, np_buffer.astype(np.int16))
-    return resampled_file
+import ffmpeg
+import whisper_live.utils as utils
 
 
 class Client:
     """
-    Handles audio recording, streaming, and communication with a server using WebSocket.
+    Handles communication with a server using WebSocket.
     """
     INSTANCES = {}
+    END_OF_AUDIO = "END_OF_AUDIO"
 
     def __init__(
-        self, host=None, port=None, is_multilingual=False, lang=None, translate=False
+        self,
+        host=None,
+        port=None,
+        lang=None,
+        translate=False,
+        model="small",
+        srt_file_path="output.srt",
+        use_vad=True,
+        log_transcription=True,
+        max_clients=4,
+        max_connection_time=600,
     ):
         """
         Initializes a Client instance for audio recording and streaming to a server.
@@ -62,38 +44,30 @@ class Client:
         Args:
             host (str): The hostname or IP address of the server.
             port (int): The port number for the WebSocket server.
-            is_multilingual (bool, optional): Specifies if multilingual transcription is enabled. Default is False.
-            lang (str, optional): The selected language for transcription when multilingual is disabled. Default is None.
+            lang (str, optional): The selected language for transcription. Default is None.
             translate (bool, optional): Specifies if the task is translation. Default is False.
         """
-        self.chunk = 1024
-        self.format = pyaudio.paInt16
-        self.channels = 1
-        self.rate = 16000
-        self.record_seconds = 60000
         self.recording = False
-        self.multilingual = False
-        self.language = None
         self.task = "transcribe"
         self.uid = str(uuid.uuid4())
         self.waiting = False
-        self.last_response_recieved = None
+        self.last_response_received = None
         self.disconnect_if_no_response_for = 15
-        self.multilingual = is_multilingual
-        self.language = lang if is_multilingual else "en"
+        self.language = lang
+        self.model = model
+        self.server_error = False
+        self.srt_file_path = srt_file_path
+        self.use_vad = use_vad
+        self.last_segment = None
+        self.last_received_segment = None
+        self.log_transcription = log_transcription
+        self.max_clients = max_clients
+        self.max_connection_time = max_connection_time
+
         if translate:
             self.task = "translate"
 
-        self.timestamp_offset = 0.0
         self.audio_bytes = None
-        self.p = pyaudio.PyAudio()
-        self.stream = self.p.open(
-            format=self.format,
-            channels=self.channels,
-            rate=self.rate,
-            input=True,
-            frames_per_buffer=self.chunk,
-        )
 
         if host is not None and port is not None:
             socket_url = f"ws://{host}:{port}"
@@ -117,13 +91,48 @@ class Client:
         self.ws_thread.setDaemon(True)
         self.ws_thread.start()
 
-        self.frames = b""
+        self.transcript = []
         print("[INFO]: * recording")
 
+    def handle_status_messages(self, message_data):
+        """Handles server status messages."""
+        status = message_data["status"]
+        if status == "WAIT":
+            self.waiting = True
+            print(f"[INFO]: Server is full. Estimated wait time {round(message_data['message'])} minutes.")
+        elif status == "ERROR":
+            print(f"Message from Server: {message_data['message']}")
+            self.server_error = True
+        elif status == "WARNING":
+            print(f"Message from Server: {message_data['message']}")
+
+    def process_segments(self, segments):
+        """Processes transcript segments."""
+        text = []
+        for i, seg in enumerate(segments):
+            if not text or text[-1] != seg["text"]:
+                text.append(seg["text"])
+                if i == len(segments) - 1 and not seg.get("completed", False):
+                    self.last_segment = seg
+                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)
+        # update last received segment and last valid response time
+        if self.last_received_segment is None or self.last_received_segment != segments[-1]["text"]:
+            self.last_response_received = time.time()
+            self.last_received_segment = segments[-1]["text"]
+
+        if self.log_transcription:
+            # Truncate to last 3 entries for brevity.
+            text = text[-3:]
+            utils.clear_screen()
+            utils.print_transcript(text)
+
     def on_message(self, ws, message):
         """
         Callback function called when a message is received from the server.
-        
+
         It updates various attributes of the client based on the received message, including
         recording status, language detection, and server messages. If a disconnect message
         is received, it sets the recording status to False.
@@ -133,25 +142,25 @@ class Client:
             message (str): The received message from the server.
 
         """
-        self.last_response_recieved = time.time()
         message = json.loads(message)
 
         if self.uid != message.get("uid"):
             print("[ERROR]: invalid client uid")
             return
 
-        if "status" in message.keys() and message["status"] == "WAIT":
-            self.waiting = True
-            print(
-                f"[INFO]:Server is full. Estimated wait time {round(message['message'])} minutes."
-            )
+        if "status" in message.keys():
+            self.handle_status_messages(message)
+            return
 
         if "message" in message.keys() and message["message"] == "DISCONNECT":
-            print("[INFO]: Server overtime disconnected.")
+            print("[INFO]: Server disconnected due to overtime.")
             self.recording = False
 
         if "message" in message.keys() and message["message"] == "SERVER_READY":
+            self.last_response_received = time.time()
             self.recording = True
+            self.server_backend = message["backend"]
+            print(f"[INFO]: Server Running with backend {self.server_backend}")
             return
 
         if "language" in message.keys():
@@ -162,78 +171,45 @@ class Client:
             )
             return
 
-        if "segments" not in message.keys():
-            return
-
-        message = message["segments"]
-        text = []
-        if len(message):
-            for seg in message:
-                if text and text[-1] == seg["text"]:
-                    # already got it
-                    continue
-                text.append(seg["text"])
-        # keep only last 3
-        if len(text) > 3:
-            text = text[-3:]
-        wrapper = textwrap.TextWrapper(width=60)
-        word_list = wrapper.wrap(text="".join(text))
-        # Print each line.
-        if os.name == "nt":
-            os.system("cls")
-        else:
-            os.system("clear")
-        for element in word_list:
-            print(element)
+        if "segments" in message.keys():
+            self.process_segments(message["segments"])
 
     def on_error(self, ws, error):
-        print(error)
+        print(f"[ERROR] WebSocket Error: {error}")
+        self.server_error = True
+        self.error_message = error
 
     def on_close(self, ws, close_status_code, close_msg):
         print(f"[INFO]: Websocket connection closed: {close_status_code}: {close_msg}")
+        self.recording = False
+        self.waiting = False
 
     def on_open(self, ws):
         """
         Callback function called when the WebSocket connection is successfully opened.
-        
-        Sends an initial configuration message to the server, including client UID, multilingual mode,
+
+        Sends an initial configuration message to the server, including client UID,
         language selection, and task type.
 
         Args:
             ws (websocket.WebSocketApp): The WebSocket client instance.
 
         """
-        print(self.multilingual, self.language, self.task)
-
         print("[INFO]: Opened connection")
         ws.send(
             json.dumps(
                 {
                     "uid": self.uid,
-                    "multilingual": self.multilingual,
                     "language": self.language,
                     "task": self.task,
+                    "model": self.model,
+                    "use_vad": self.use_vad,
+                    "max_clients": self.max_clients,
+                    "max_connection_time": self.max_connection_time,
                 }
             )
         )
 
-    @staticmethod
-    def bytes_to_float_array(audio_bytes):
-        """
-        Convert audio data from bytes to a NumPy float array.
-        
-        It assumes that the audio data is in 16-bit PCM format. The audio data is normalized to 
-        have values between -1 and 1.
-
-        Args:
-            audio_bytes (bytes): Audio data in bytes.
-
-        Returns:
-            np.ndarray: A NumPy array containing the audio data as float values normalized between -1 and 1.
-        """
-        raw_data = np.frombuffer(buffer=audio_bytes, dtype=np.int16)
-        return raw_data.astype(np.float32) / 32768.0
-
     def send_packet_to_server(self, message):
         """
         Send an audio packet to the server using WebSocket.
@@ -247,62 +223,11 @@ class Client:
         except Exception as e:
             print(e)
 
-    def play_file(self, filename):
-        """
-        Play an audio file and send it to the server for processing.
-        
-        Reads an audio file, plays it through the audio output, and simultaneously sends
-        the audio data to the server for processing. It uses PyAudio to create an audio 
-        stream for playback. The audio data is read from the file in chunks, converted to 
-        floating-point format, and sent to the server using WebSocket communication.
-        This method is typically used when you want to process pre-recorded audio and send it
-        to the server in real-time.
-
-        Args:
-            filename (str): The path to the audio file to be played and sent to the server.
-        """
-        
-        # read audio and create pyaudio stream
-        with wave.open(filename, "rb") as wavfile:
-            self.stream = self.p.open(
-                format=self.p.get_format_from_width(wavfile.getsampwidth()),
-                channels=wavfile.getnchannels(),
-                rate=wavfile.getframerate(),
-                input=True,
-                output=True,
-                frames_per_buffer=self.chunk,
-            )
-            try:
-                while self.recording:
-                    data = wavfile.readframes(self.chunk)
-                    if data == b"":
-                        break
-
-                    audio_array = self.bytes_to_float_array(data)
-                    self.send_packet_to_server(audio_array.tobytes())
-                    self.stream.write(data)
-
-                wavfile.close()
-
-                assert self.last_response_recieved
-                while time.time() - self.last_response_recieved < self.disconnect_if_no_response_for:
-                    continue
-                self.stream.close()
-                self.close_websocket()
-
-            except KeyboardInterrupt:
-                wavfile.close()
-                self.stream.stop_stream()
-                self.stream.close()
-                self.p.terminate()
-                self.close_websocket()
-                print("[INFO]: Keyboard interrupt.")
-
     def close_websocket(self):
         """
         Close the WebSocket connection and join the WebSocket thread.
 
-        First attempts to close the WebSocket connection using `self.client_socket.close()`. After 
+        First attempts to close the WebSocket connection using `self.client_socket.close()`. After
         closing the connection, it joins the WebSocket thread to ensure proper termination.
 
         """
@@ -325,11 +250,337 @@ class Client:
         """
         return self.client_socket
 
+    def write_srt_file(self, output_path="output.srt"):
+        """
+        Writes out the transcript in .srt format.
+
+        Args:
+            message (output_path, optional): The path to the target file.  Default is "output.srt".
+
+        """
+        if self.server_backend == "faster_whisper":
+            if not self.transcript and self.last_segment is not None:
+                self.transcript.append(self.last_segment)
+            elif self.last_segment and self.transcript[-1]["text"] != self.last_segment["text"]:
+                self.transcript.append(self.last_segment)
+            utils.create_srt_file(self.transcript, output_path)
+
+    def wait_before_disconnect(self):
+        """Waits a bit before disconnecting in order to process pending responses."""
+        assert self.last_response_received
+        while time.time() - self.last_response_received < self.disconnect_if_no_response_for:
+            continue
+
+
+class TranscriptionTeeClient:
+    """
+    Client for handling audio recording, streaming, and transcription tasks via one or more
+    WebSocket connections.
+
+    Acts as a high-level client for audio transcription tasks using a WebSocket connection. It can be used
+    to send audio data for transcription to one or more servers, and receive transcribed text segments.
+    Args:
+        clients (list): one or more previously initialized Client instances
+
+    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"):
+        self.clients = clients
+        if not self.clients:
+            raise Exception("At least one client is required.")
+        self.chunk = 4096
+        self.format = pyaudio.paInt16
+        self.channels = 1
+        self.rate = 16000
+        self.record_seconds = 60000
+        self.save_output_recording = save_output_recording
+        self.output_recording_filename = output_recording_filename
+        self.frames = b""
+        self.p = pyaudio.PyAudio()
+        try:
+            self.stream = self.p.open(
+                format=self.format,
+                channels=self.channels,
+                rate=self.rate,
+                input=True,
+                frames_per_buffer=self.chunk,
+            )
+        except OSError as error:
+            print(f"[WARN]: Unable to access microphone. {error}")
+            self.stream = None
+
+    def __call__(self, audio=None, rtsp_url=None, hls_url=None, save_file=None):
+        """
+        Start the transcription process.
+
+        Initiates the transcription process by connecting to the server via a WebSocket. It waits for the server
+        to be ready to receive audio data and then sends audio for transcription. If an audio file is provided, it
+        will be played and streamed to the server; otherwise, it will perform live recording.
+
+        Args:
+            audio (str, optional): Path to an audio file for transcription. Default is None, which triggers live recording.
+
+        """
+        assert sum(
+            source is not None for source in [audio, rtsp_url, hls_url]
+        ) <= 1, 'You must provide only one selected source'
+
+        print("[INFO]: Waiting for server ready ...")
+        for client in self.clients:
+            while not client.recording:
+                if client.waiting or client.server_error:
+                    self.close_all_clients()
+                    return
+
+        print("[INFO]: Server Ready!")
+        if hls_url is not None:
+            self.process_hls_stream(hls_url, save_file)
+        elif audio is not None:
+            resampled_file = utils.resample(audio)
+            self.play_file(resampled_file)
+        elif rtsp_url is not None:
+            self.process_rtsp_stream(rtsp_url)
+        else:
+            self.record()
+
+    def close_all_clients(self):
+        """Closes all client websockets."""
+        for client in self.clients:
+            client.close_websocket()
+
+    def write_all_clients_srt(self):
+        """Writes out .srt files for all clients."""
+        for client in self.clients:
+            client.write_srt_file(client.srt_file_path)
+
+    def multicast_packet(self, packet, unconditional=False):
+        """
+        Sends an identical packet via all clients.
+
+        Args:
+            packet (bytes): The audio data packet in bytes to be sent.
+            unconditional (bool, optional): If true, send regardless of whether clients are recording.  Default is False.
+        """
+        for client in self.clients:
+            if (unconditional or client.recording):
+                client.send_packet_to_server(packet)
+
+    def play_file(self, filename):
+        """
+        Play an audio file and send it to the server for processing.
+
+        Reads an audio file, plays it through the audio output, and simultaneously sends
+        the audio data to the server for processing. It uses PyAudio to create an audio
+        stream for playback. The audio data is read from the file in chunks, converted to
+        floating-point format, and sent to the server using WebSocket communication.
+        This method is typically used when you want to process pre-recorded audio and send it
+        to the server in real-time.
+
+        Args:
+            filename (str): The path to the audio file to be played and sent to the server.
+        """
+
+        # read audio and create pyaudio stream
+        with wave.open(filename, "rb") as wavfile:
+            self.stream = self.p.open(
+                format=self.p.get_format_from_width(wavfile.getsampwidth()),
+                channels=wavfile.getnchannels(),
+                rate=wavfile.getframerate(),
+                input=True,
+                output=True,
+                frames_per_buffer=self.chunk,
+            )
+            try:
+                while any(client.recording for client in self.clients):
+                    data = wavfile.readframes(self.chunk)
+                    if data == b"":
+                        break
+
+                    audio_array = self.bytes_to_float_array(data)
+                    self.multicast_packet(audio_array.tobytes())
+                    self.stream.write(data)
+
+                wavfile.close()
+
+                for client in self.clients:
+                    client.wait_before_disconnect()
+                self.multicast_packet(Client.END_OF_AUDIO.encode('utf-8'), True)
+                self.write_all_clients_srt()
+                self.stream.close()
+                self.close_all_clients()
+
+            except KeyboardInterrupt:
+                wavfile.close()
+                self.stream.stop_stream()
+                self.stream.close()
+                self.p.terminate()
+                self.close_all_clients()
+                self.write_all_clients_srt()
+                print("[INFO]: Keyboard interrupt.")
+
+    def process_rtsp_stream(self, rtsp_url):
+        """
+        Connect to an RTSP source, process the audio stream, and send it for trascription.
+
+        Args:
+            rtsp_url (str): The URL of the RTSP stream source.
+        """
+        process = self.get_rtsp_ffmpeg_process(rtsp_url)
+        self.handle_ffmpeg_process(process, stream_type='RTSP')
+
+    def process_hls_stream(self, hls_url, save_file):
+        """
+        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.
+        """
+        process = self.get_hls_ffmpeg_process(hls_url, save_file)
+        self.handle_ffmpeg_process(process, stream_type='HLS')
+
+    def handle_ffmpeg_process(self, process, stream_type):
+        print(f"[INFO]: Connecting to {stream_type} stream...")
+        stderr_thread = threading.Thread(target=self.consume_stderr, args=(process,))
+        stderr_thread.start()
+        try:
+            # Process the stream
+            while True:
+                in_bytes = process.stdout.read(self.chunk * 2)  # 2 bytes per sample
+                if not in_bytes:
+                    break
+                audio_array = self.bytes_to_float_array(in_bytes)
+                self.multicast_packet(audio_array.tobytes())
+
+        except Exception as e:
+            print(f"[ERROR]: Failed to connect to {stream_type} stream: {e}")
+        finally:
+            self.close_all_clients()
+            self.write_all_clients_srt()
+            if process:
+                process.kill()
+
+        print(f"[INFO]: {stream_type} stream processing finished.")
+
+    def get_rtsp_ffmpeg_process(self, rtsp_url):
+        return (
+            ffmpeg
+            .input(rtsp_url, threads=0)
+            .output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
+            .run_async(pipe_stdout=True, pipe_stderr=True)
+        )
+
+    def get_hls_ffmpeg_process(self, hls_url, save_file):
+        if save_file is None:
+            process = (
+                ffmpeg
+                .input(hls_url, threads=0)
+                .output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
+                .run_async(pipe_stdout=True, pipe_stderr=True)
+            )
+        else:
+            input = ffmpeg.input(hls_url, threads=0)
+            output_file = input.output(save_file, acodec='copy', vcodec='copy').global_args('-loglevel', 'quiet')
+            output_std = input.output('-', format='s16le', acodec='pcm_s16le', ac=1, ar=self.rate)
+            process = (
+                ffmpeg.merge_outputs(output_file, output_std)
+                .run_async(pipe_stdout=True, pipe_stderr=True)
+            )
+
+        return process
+
+    def consume_stderr(self, process):
+        """
+        Consume and log the stderr output of a process in a separate thread.
+
+        Args:
+            process (subprocess.Popen): The process whose stderr output will be logged.
+        """
+        for line in iter(process.stderr.readline, b""):
+            logging.debug(f'[STDERR]: {line.decode()}')
+
+    def save_chunk(self, n_audio_file):
+        """
+        Saves the current audio frames to a WAV file in a separate thread.
+
+        Args:
+        n_audio_file (int): The index of the audio file which determines the filename.
+                            This helps in maintaining the order and uniqueness of each chunk.
+        """
+        t = threading.Thread(
+            target=self.write_audio_frames_to_file,
+            args=(self.frames[:], f"chunks/{n_audio_file}.wav",),
+        )
+        t.start()
+
+    def finalize_recording(self, n_audio_file):
+        """
+        Finalizes the recording process by saving any remaining audio frames,
+        closing the audio stream, and terminating the process.
+
+        Args:
+        n_audio_file (int): The file index to be used if there are remaining audio frames to be saved.
+                            This index is incremented before use if the last chunk is saved.
+        """
+        if self.save_output_recording and len(self.frames):
+            self.write_audio_frames_to_file(
+                self.frames[:], f"chunks/{n_audio_file}.wav"
+            )
+            n_audio_file += 1
+        self.stream.stop_stream()
+        self.stream.close()
+        self.p.terminate()
+        self.close_all_clients()
+        if self.save_output_recording:
+            self.write_output_recording(n_audio_file)
+        self.write_all_clients_srt()
+
+    def record(self):
+        """
+        Record audio data from the input stream and save it to a WAV file.
+
+        Continuously records audio data from the input stream, sends it to the server via a WebSocket
+        connection, and simultaneously saves it to multiple WAV files in chunks. It stops recording when
+        the `RECORD_SECONDS` duration is reached or when the `RECORDING` flag is set to `False`.
+
+        Audio data is saved in chunks to the "chunks" directory. Each chunk is saved as a separate WAV file.
+        The recording will continue until the specified duration is reached or until the `RECORDING` flag is set to `False`.
+        The recording process can be interrupted by sending a KeyboardInterrupt (e.g., pressing Ctrl+C). After recording,
+        the method combines all the saved audio chunks into the specified `out_file`.
+        """
+        n_audio_file = 0
+        if self.save_output_recording:
+            if os.path.exists("chunks"):
+                shutil.rmtree("chunks")
+            os.makedirs("chunks")
+        try:
+            for _ in range(0, int(self.rate / self.chunk * self.record_seconds)):
+                if not any(client.recording for client in self.clients):
+                    break
+                data = self.stream.read(self.chunk, exception_on_overflow=False)
+                self.frames += data
+
+                audio_array = self.bytes_to_float_array(data)
+
+                self.multicast_packet(audio_array.tobytes())
+
+                # save frames if more than a minute
+                if len(self.frames) > 60 * self.rate:
+                    if self.save_output_recording:
+                        self.save_chunk(n_audio_file)
+                        n_audio_file += 1
+                    self.frames = b""
+            self.write_all_clients_srt()
+
+        except KeyboardInterrupt:
+            self.finalize_recording(n_audio_file)
+
     def write_audio_frames_to_file(self, frames, file_name):
         """
         Write audio frames to a WAV file.
 
-        The WAV file is created or overwritten with the specified name. The audio frames should be 
+        The WAV file is created or overwritten with the specified name. The audio frames should be
         in the correct format and match the specified channel, sample width, and sample rate.
 
         Args:
@@ -344,68 +595,11 @@ class Client:
             wavfile.setframerate(self.rate)
             wavfile.writeframes(frames)
 
-    def record(self, out_file="output_recording.wav"):
-        """
-        Record audio data from the input stream and save it to a WAV file.
-
-        Continuously records audio data from the input stream, sends it to the server via a WebSocket
-        connection, and simultaneously saves it to multiple WAV files in chunks. It stops recording when
-        the `RECORD_SECONDS` duration is reached or when the `RECORDING` flag is set to `False`.
-
-        Audio data is saved in chunks to the "chunks" directory. Each chunk is saved as a separate WAV file.
-        The recording will continue until the specified duration is reached or until the `RECORDING` flag is set to `False`.
-        The recording process can be interrupted by sending a KeyboardInterrupt (e.g., pressing Ctrl+C). After recording, 
-        the method combines all the saved audio chunks into the specified `out_file`.
-
-        Args:
-            out_file (str, optional): The name of the output WAV file to save the entire recording. Default is "output_recording.wav".
-
-        """
-        n_audio_file = 0
-        if not os.path.exists("chunks"):
-            os.makedirs("chunks", exist_ok=True)
-        try:
-            for _ in range(0, int(self.rate / self.chunk * self.record_seconds)):
-                if not self.recording:
-                    break
-                data = self.stream.read(self.chunk)
-                self.frames += data
-
-                audio_array = Client.bytes_to_float_array(data)
-
-                self.send_packet_to_server(audio_array.tobytes())
-
-                # save frames if more than a minute
-                if len(self.frames) > 60 * self.rate:
-                    t = threading.Thread(
-                        target=self.write_audio_frames_to_file,
-                        args=(
-                            self.frames[:],
-                            f"chunks/{n_audio_file}.wav",
-                        ),
-                    )
-                    t.start()
-                    n_audio_file += 1
-                    self.frames = b""
-
-        except KeyboardInterrupt:
-            if len(self.frames):
-                self.write_audio_frames_to_file(
-                    self.frames[:], f"chunks/{n_audio_file}.wav"
-                )
-                n_audio_file += 1
-            self.stream.stop_stream()
-            self.stream.close()
-            self.p.terminate()
-            self.close_websocket()
-
-            self.write_output_recording(n_audio_file, out_file)
-
-    def write_output_recording(self, n_audio_file, out_file):
+    def write_output_recording(self, n_audio_file):
         """
         Combine and save recorded audio chunks into a single WAV file.
-        
-        The individual audio chunk files are expected to be located in the "chunks" directory. Reads each chunk 
+
+        The individual audio chunk files are expected to be located in the "chunks" directory. Reads each chunk
         file, appends its audio data to the final recording, and then deletes the chunk file. After combining
         and saving, the final recording is stored in the specified `out_file`.
 
@@ -420,7 +614,7 @@ class Client:
             for i in range(n_audio_file)
             if os.path.exists(f"chunks/{i}.wav")
         ]
-        with wave.open(out_file, "wb") as wavfile:
+        with wave.open(self.output_recording_filename, "wb") as wavfile:
             wavfile: wave.Wave_write
             wavfile.setnchannels(self.channels)
             wavfile.setsampwidth(2)
@@ -435,11 +629,31 @@ class Client:
                 # remove this file
                 os.remove(in_file)
         wavfile.close()
+        # clean up temporary directory to store chunks
+        if os.path.exists("chunks"):
+            shutil.rmtree("chunks")
+
+    @staticmethod
+    def bytes_to_float_array(audio_bytes):
+        """
+        Convert audio data from bytes to a NumPy float array.
+
+        It assumes that the audio data is in 16-bit PCM format. The audio data is normalized to
+        have values between -1 and 1.
+
+        Args:
+            audio_bytes (bytes): Audio data in bytes.
+
+        Returns:
+            np.ndarray: A NumPy array containing the audio data as float values normalized between -1 and 1.
+        """
+        raw_data = np.frombuffer(buffer=audio_bytes, dtype=np.int16)
+        return raw_data.astype(np.float32) / 32768.0
 
 
-class TranscriptionClient:
+class TranscriptionClient(TranscriptionTeeClient):
     """
-    Client for handling audio transcription tasks via a WebSocket connection.
+    Client for handling audio transcription tasks via a single WebSocket connection.
 
     Acts as a high-level client for audio transcription tasks using a WebSocket connection. It can be used
     to send audio data for transcription to a server and receive transcribed text segments.
@@ -447,9 +661,11 @@ class TranscriptionClient:
     Args:
         host (str): The hostname or IP address of the server.
         port (int): The port number to connect to on the server.
-        is_multilingual (bool, optional): Indicates whether the transcription should support multiple languages (default is False).
-        lang (str, optional): The primary language for transcription (used if `is_multilingual` is False). Default is None, which defaults to English ('en').
+        lang (str, optional): The primary language for transcription. Default is None, which defaults to English ('en').
         translate (bool, optional): Indicates whether translation tasks are required (default is False).
+        save_output_recording (bool, optional): Indicates whether to save recording from microphone.
+        output_recording_filename (str, optional): File to save the output recording.
+        output_transcription_path (str, optional): File to save the output transcription.
 
     Attributes:
         client (Client): An instance of the underlying Client class responsible for handling the WebSocket connection.
@@ -457,34 +673,38 @@ class TranscriptionClient:
     Example:
         To create a TranscriptionClient and start transcription on microphone audio:
         ```python
-        transcription_client = TranscriptionClient(host="localhost", port=9090, is_multilingual=True)
+        transcription_client = TranscriptionClient(host="localhost", port=9090)
         transcription_client()
         ```
     """
-    def __init__(self, host, port, is_multilingual=False, lang=None, translate=False):
-        self.client = Client(host, port, is_multilingual, lang, translate)
+    def __init__(
+        self,
+        host,
+        port,
+        lang=None,
+        translate=False,
+        model="small",
+        use_vad=True,
+        save_output_recording=False,
+        output_recording_filename="./output_recording.wav",
+        output_transcription_path="./output.srt",
+        log_transcription=True,
+        max_clients=4,
+        max_connection_time=600,
+    ):
+        self.client = Client(
+            host, port, lang, translate, model, srt_file_path=output_transcription_path,
+            use_vad=use_vad, log_transcription=log_transcription, max_clients=max_clients,
+            max_connection_time=max_connection_time
+        )
 
-    def __call__(self, audio=None):
-        """
-        Start the transcription process.
-
-        Initiates the transcription process by connecting to the server via a WebSocket. It waits for the server
-        to be ready to receive audio data and then sends audio for transcription. If an audio file is provided, it 
-        will be played and streamed to the server; otherwise, it will perform live recording.
-
-        Args:
-            audio (str, optional): Path to an audio file for transcription. Default is None, which triggers live recording.
-                   
-        """
-        print("[INFO]: Waiting for server ready ...")
-        while not self.client.recording:
-            if self.client.waiting:
-                self.client.close_websocket()
-                return
-            pass
-        print("[INFO]: Server Ready!")
-        if audio is not None:
-            resampled_file = resample(audio)
-            self.client.play_file(resampled_file)
-        else:
-            self.client.record()
+        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"):
+            raise ValueError(f"Please provide a valid `output_transcription_path`: {output_transcription_path}. The file extension should be `.srt`.")
+        TranscriptionTeeClient.__init__(
+            self,
+            [self.client],
+            save_output_recording=save_output_recording,
+            output_recording_filename=output_recording_filename
+        )

whisper_live/tensorrt_utils.py

@@ -0,0 +1,365 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import logging
+import os
+from collections import defaultdict
+from functools import lru_cache
+from pathlib import Path
+from subprocess import CalledProcessError, run
+from typing import Dict, Iterable, List, Optional, TextIO, Tuple, Union
+
+import kaldialign
+import numpy as np
+import soundfile
+import torch
+import torch.nn.functional as F
+
+Pathlike = Union[str, Path]
+
+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 load_audio(file: str, sr: int = SAMPLE_RATE):
+    """
+    Open an audio file and read as mono waveform, resampling as necessary
+
+    Parameters
+    ----------
+    file: str
+        The audio file to open
+
+    sr: int
+        The sample rate to resample the audio if necessary
+
+    Returns
+    -------
+    A NumPy array containing the audio waveform, in float32 dtype.
+    """
+
+    # This launches a subprocess to decode audio while down-mixing
+    # and resampling as necessary.  Requires the ffmpeg CLI in PATH.
+    # fmt: off
+    cmd = [
+        "ffmpeg", "-nostdin", "-threads", "0", "-i", file, "-f", "s16le", "-ac",
+        "1", "-acodec", "pcm_s16le", "-ar",
+        str(sr), "-"
+    ]
+    # fmt: on
+    try:
+        out = run(cmd, capture_output=True, check=True).stdout
+    except CalledProcessError as e:
+        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
+
+    return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+
+
+def load_audio_wav_format(wav_path):
+    # make sure audio in .wav format
+    assert wav_path.endswith(
+        '.wav'), f"Only support .wav format, but got {wav_path}"
+    waveform, sample_rate = soundfile.read(wav_path)
+    assert sample_rate == 16000, f"Only support 16k sample rate, but got {sample_rate}"
+    return waveform, sample_rate
+
+
+def pad_or_trim(array, length: int = N_SAMPLES, *, axis: int = -1):
+    """
+    Pad or trim the audio array to N_SAMPLES, as expected by the encoder.
+    """
+    if torch.is_tensor(array):
+        if array.shape[axis] > length:
+            array = array.index_select(dim=axis,
+                                       index=torch.arange(length,
+                                                          device=array.device))
+
+        if array.shape[axis] < length:
+            pad_widths = [(0, 0)] * array.ndim
+            pad_widths[axis] = (0, length - array.shape[axis])
+            array = F.pad(array,
+                          [pad for sizes in pad_widths[::-1] for pad in sizes])
+    else:
+        if array.shape[axis] > length:
+            array = array.take(indices=range(length), axis=axis)
+
+        if array.shape[axis] < length:
+            pad_widths = [(0, 0)] * array.ndim
+            pad_widths[axis] = (0, length - array.shape[axis])
+            array = np.pad(array, pad_widths)
+
+    return array
+
+
+@lru_cache(maxsize=None)
+def mel_filters(device,
+                n_mels: int,
+                mel_filters_dir: str = None) -> torch.Tensor:
+    """
+    load the mel filterbank matrix for projecting STFT into a Mel spectrogram.
+    Allows decoupling librosa dependency; saved using:
+
+        np.savez_compressed(
+            "mel_filters.npz",
+            mel_80=librosa.filters.mel(sr=16000, n_fft=400, n_mels=80),
+        )
+    """
+    assert n_mels in {80, 128}, f"Unsupported n_mels: {n_mels}"
+    if mel_filters_dir is None:
+        mel_filters_path = os.path.join(os.path.dirname(__file__), "assets",
+                                        "mel_filters.npz")
+    else:
+        mel_filters_path = os.path.join(mel_filters_dir, "mel_filters.npz")
+    with np.load(mel_filters_path) as f:
+        return torch.from_numpy(f[f"mel_{n_mels}"]).to(device)
+
+
+def log_mel_spectrogram(
+    audio: Union[str, np.ndarray, torch.Tensor],
+    n_mels: int,
+    padding: int = 0,
+    device: Optional[Union[str, torch.device]] = None,
+    return_duration: bool = False,
+    mel_filters_dir: str = None,
+):
+    """
+    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 device is not None:
+        audio = audio.to(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
+
+    filters = mel_filters(audio.device, n_mels, mel_filters_dir)
+    mel_spec = 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 store_transcripts(filename: Pathlike, texts: Iterable[Tuple[str, str,
+                                                                str]]) -> None:
+    """Save predicted results and reference transcripts to a file.
+    https://github.com/k2-fsa/icefall/blob/master/icefall/utils.py
+    Args:
+      filename:
+        File to save the results to.
+      texts:
+        An iterable of tuples. The first element is the cur_id, the second is
+        the reference transcript and the third element is the predicted result.
+    Returns:
+      Return None.
+    """
+    with open(filename, "w") as f:
+        for cut_id, ref, hyp in texts:
+            print(f"{cut_id}:\tref={ref}", file=f)
+            print(f"{cut_id}:\thyp={hyp}", file=f)
+
+
+def write_error_stats(                                              # noqa: C901
+    f: TextIO,
+    test_set_name: str,
+    results: List[Tuple[str, str]],
+    enable_log: bool = True,
+) -> float:
+    """Write statistics based on predicted results and reference transcripts.
+    https://github.com/k2-fsa/icefall/blob/master/icefall/utils.py
+    It will write the following to the given file:
+
+        - WER
+        - number of insertions, deletions, substitutions, corrects and total
+          reference words. For example::
+
+              Errors: 23 insertions, 57 deletions, 212 substitutions, over 2606
+              reference words (2337 correct)
+
+        - The difference between the reference transcript and predicted result.
+          An instance is given below::
+
+            THE ASSOCIATION OF (EDISON->ADDISON) ILLUMINATING COMPANIES
+
+          The above example shows that the reference word is `EDISON`,
+          but it is predicted to `ADDISON` (a substitution error).
+
+          Another example is::
+
+            FOR THE FIRST DAY (SIR->*) I THINK
+
+          The reference word `SIR` is missing in the predicted
+          results (a deletion error).
+      results:
+        An iterable of tuples. The first element is the cur_id, the second is
+        the reference transcript and the third element is the predicted result.
+      enable_log:
+        If True, also print detailed WER to the console.
+        Otherwise, it is written only to the given file.
+    Returns:
+      Return None.
+    """
+    subs: Dict[Tuple[str, str], int] = defaultdict(int)
+    ins: Dict[str, int] = defaultdict(int)
+    dels: Dict[str, int] = defaultdict(int)
+
+    # `words` stores counts per word, as follows:
+    #   corr, ref_sub, hyp_sub, ins, dels
+    words: Dict[str, List[int]] = defaultdict(lambda: [0, 0, 0, 0, 0])
+    num_corr = 0
+    ERR = "*"
+    for cut_id, ref, hyp in results:
+        ali = kaldialign.align(ref, hyp, ERR)
+        for ref_word, hyp_word in ali:
+            if ref_word == ERR:
+                ins[hyp_word] += 1
+                words[hyp_word][3] += 1
+            elif hyp_word == ERR:
+                dels[ref_word] += 1
+                words[ref_word][4] += 1
+            elif hyp_word != ref_word:
+                subs[(ref_word, hyp_word)] += 1
+                words[ref_word][1] += 1
+                words[hyp_word][2] += 1
+            else:
+                words[ref_word][0] += 1
+                num_corr += 1
+    ref_len = sum([len(r) for _, r, _ in results])
+    sub_errs = sum(subs.values())
+    ins_errs = sum(ins.values())
+    del_errs = sum(dels.values())
+    tot_errs = sub_errs + ins_errs + del_errs
+    tot_err_rate = "%.2f" % (100.0 * tot_errs / ref_len)
+
+    if enable_log:
+        logging.info(f"[{test_set_name}] %WER {tot_errs / ref_len:.2%} "
+                     f"[{tot_errs} / {ref_len}, {ins_errs} ins, "
+                     f"{del_errs} del, {sub_errs} sub ]")
+
+    print(f"%WER = {tot_err_rate}", file=f)
+    print(
+        f"Errors: {ins_errs} insertions, {del_errs} deletions, "
+        f"{sub_errs} substitutions, over {ref_len} reference "
+        f"words ({num_corr} correct)",
+        file=f,
+    )
+    print(
+        "Search below for sections starting with PER-UTT DETAILS:, "
+        "SUBSTITUTIONS:, DELETIONS:, INSERTIONS:, PER-WORD STATS:",
+        file=f,
+    )
+
+    print("", file=f)
+    print("PER-UTT DETAILS: corr or (ref->hyp)  ", file=f)
+    for cut_id, ref, hyp in results:
+        ali = kaldialign.align(ref, hyp, ERR)
+        combine_successive_errors = True
+        if combine_successive_errors:
+            ali = [[[x], [y]] for x, y in ali]
+            for i in range(len(ali) - 1):
+                if ali[i][0] != ali[i][1] and ali[i + 1][0] != ali[i + 1][1]:
+                    ali[i + 1][0] = ali[i][0] + ali[i + 1][0]
+                    ali[i + 1][1] = ali[i][1] + ali[i + 1][1]
+                    ali[i] = [[], []]
+            ali = [[
+                list(filter(lambda a: a != ERR, x)),
+                list(filter(lambda a: a != ERR, y)),
+            ] for x, y in ali]
+            ali = list(filter(lambda x: x != [[], []], ali))
+            ali = [[
+                ERR if x == [] else " ".join(x),
+                ERR if y == [] else " ".join(y),
+            ] for x, y in ali]
+
+        print(
+            f"{cut_id}:\t" + " ".join((ref_word if ref_word == hyp_word else
+                                       f"({ref_word}->{hyp_word})"
+                                       for ref_word, hyp_word in ali)),
+            file=f,
+        )
+
+    print("", file=f)
+    print("SUBSTITUTIONS: count ref -> hyp", file=f)
+
+    for count, (ref, hyp) in sorted([(v, k) for k, v in subs.items()],
+                                    reverse=True):
+        print(f"{count}   {ref} -> {hyp}", file=f)
+
+    print("", file=f)
+    print("DELETIONS: count ref", file=f)
+    for count, ref in sorted([(v, k) for k, v in dels.items()], reverse=True):
+        print(f"{count}   {ref}", file=f)
+
+    print("", file=f)
+    print("INSERTIONS: count hyp", file=f)
+    for count, hyp in sorted([(v, k) for k, v in ins.items()], reverse=True):
+        print(f"{count}   {hyp}", file=f)
+
+    print("", file=f)
+    print("PER-WORD STATS: word  corr tot_errs count_in_ref count_in_hyp",
+          file=f)
+    for _, word, counts in sorted([(sum(v[1:]), k, v)
+                                   for k, v in words.items()],
+                                  reverse=True):
+        (corr, ref_sub, hyp_sub, ins, dels) = counts
+        tot_errs = ref_sub + hyp_sub + ins + dels
+        ref_count = corr + ref_sub + dels
+        hyp_count = corr + hyp_sub + ins
+
+        print(f"{word}   {corr} {tot_errs} {ref_count} {hyp_count}", file=f)
+    return float(tot_err_rate)

whisper_live/transcriber_tensorrt.py

@@ -0,0 +1,421 @@
+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.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
+
+
+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"):
+        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.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.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,
+            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)
+
+        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)
+        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",
+            ):
+        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)
+        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

@@ -0,0 +1,71 @@
+import os
+import textwrap
+import scipy
+import ffmpeg
+import numpy as np
+
+
+def clear_screen():
+    """Clears the console screen."""
+    os.system("cls" if os.name == "nt" else "clear")
+
+
+def print_transcript(text):
+    """Prints formatted transcript text."""
+    wrapper = textwrap.TextWrapper(width=60)
+    for line in wrapper.wrap(text="".join(text)):
+        print(line)
+
+
+def format_time(s):
+    """Convert seconds (float) to SRT time format."""
+    hours = int(s // 3600)
+    minutes = int((s % 3600) // 60)
+    seconds = int(s % 60)
+    milliseconds = int((s - int(s)) * 1000)
+    return f"{hours:02}:{minutes:02}:{seconds:02},{milliseconds:03}"
+
+
+def create_srt_file(segments, output_file):
+    with open(output_file, 'w', encoding='utf-8') as srt_file:
+        segment_number = 1
+        for segment in segments:
+            start_time = format_time(float(segment['start']))
+            end_time = format_time(float(segment['end']))
+            text = segment['text']
+
+            srt_file.write(f"{segment_number}\n")
+            srt_file.write(f"{start_time} --> {end_time}\n")
+            srt_file.write(f"{text}\n\n")
+
+            segment_number += 1
+
+
+def resample(file: str, sr: int = 16000):
+    """
+    # https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/audio.py#L22
+    Open an audio file and read as mono waveform, resampling as necessary,
+    save the resampled audio
+
+    Args:
+        file (str): The audio file to open
+        sr (int): The sample rate to resample the audio if necessary
+
+    Returns:
+        resampled_file (str): The resampled audio file
+    """
+    try:
+        # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
+        # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
+        out, _ = (
+            ffmpeg.input(file, threads=0)
+            .output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=sr)
+            .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
+        )
+    except ffmpeg.Error as e:
+        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
+    np_buffer = np.frombuffer(out, dtype=np.int16)
+
+    resampled_file = f"{file.split('.')[0]}_resampled.wav"
+    scipy.io.wavfile.write(resampled_file, sr, np_buffer.astype(np.int16))
+    return resampled_file

whisper_live/vad.py

@@ -11,6 +11,7 @@ class VoiceActivityDetection():
 
     def __init__(self, force_onnx_cpu=True):
         path = self.download()
+
         opts = onnxruntime.SessionOptions()
         opts.log_severity_level = 3
 
@@ -22,7 +23,6 @@ class VoiceActivityDetection():
         else:
             self.session = onnxruntime.InferenceSession(path, providers=['CUDAExecutionProvider'], sess_options=opts)
 
-
         self.reset_states()
         self.sample_rates = [8000, 16000]
 
@@ -34,7 +34,7 @@ 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:
@@ -94,7 +94,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/v4.0/files/silero_vad.onnx"):
         target_dir = os.path.expanduser("~/.cache/whisper-live/")
 
         # Ensure the target directory exists
@@ -106,10 +106,37 @@ class VoiceActivityDetection():
         # Check if the model file already exists
         if not os.path.exists(model_filename):
             # If it doesn't exist, download the model using wget
-            print("Downloading VAD ONNX model...")
             try:
                 subprocess.run(["wget", "-O", model_filename, model_url], check=True)
             except subprocess.CalledProcessError:
                 print("Failed to download the model using wget.")
         return model_filename
 
+
+class VoiceActivityDetector:
+    def __init__(self, threshold=0.5, frame_rate=16000):
+        """
+        Initializes the VoiceActivityDetector with a voice activity detection model and a threshold.
+
+        Args:
+            threshold (float, optional): The probability threshold for detecting voice activity. Defaults to 0.5.
+        """
+        self.model = VoiceActivityDetection()
+        self.threshold = threshold
+        self.frame_rate = frame_rate
+
+    def __call__(self, audio_frame):
+        """
+        Determines if the given audio frame contains speech by comparing the detected speech probability against
+        the threshold.
+
+        Args:
+            audio_frame (np.ndarray): The audio frame to be analyzed for voice activity. It is expected to be a
+                                      NumPy array of audio samples.
+
+        Returns:
+            bool: True if the speech probability exceeds the threshold, indicating the presence of voice activity;
+                  False otherwise.
+        """
+        speech_prob = self.model(torch.from_numpy(audio_frame), self.frame_rate).item()
+        return speech_prob > self.threshold