bump version v0.1.0

Firefox remove multilingual

Audio-Transcription-Chrome/README.md

@@ -26,7 +26,6 @@ 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.

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');

Audio-Transcription-Chrome/options.js

@@ -93,17 +93,13 @@ 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,
-          model_size: option.modelSize
+          model: option.modelSize
         })
       );
     };

Audio-Transcription-Chrome/popup.html

@@ -15,112 +15,109 @@
     <input type="checkbox" id="useServerCheckbox">
     <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>
-  </div>
   <div class="dropdown-container">
     <label for="languageDropdown">Select Language:</label>
-    <select id="languageDropdown" disabled>
+    <select id="languageDropdown">
-      <option value="">Select Language</option>
+      <option value="" selected>Automatically detect</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>
       <option value="af">Afrikaans</option>
-      <option value="oc">Occitan</option>
+      <option value="sq">Albanian</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="am">Amharic</option>
-      <option value="yi">Yiddish</option>
+      <option value="ar">Arabic</option>
-      <option value="lo">Lao</option>
+      <option value="hy">Armenian</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="as">Assamese</option>
-      <option value="tt">Tatar</option>
+      <option value="az">Azerbaijani</option>
-      <option value="haw">Hawaiian</option>
-      <option value="ln">Lingala</option>
-      <option value="ha">Hausa</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">
@@ -134,11 +131,15 @@
   <div class="dropdown-container">
     <label for="modelSizeDropdown">Select Model Size:</label>
     <select id="modelSizeDropdown">
-      <option value="">Select Task</option>
+      <option value="">Select model</option>
-      <option value="tiny">Tiny</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>

Audio-Transcription-Chrome/popup.js

@@ -4,7 +4,6 @@ document.addEventListener("DOMContentLoaded", function () {
   const stopButton = document.getElementById("stopCapture");
 
   const useServerCheckbox = document.getElementById("useServerCheckbox");
-  const useMultilingualCheckbox = document.getElementById('useMultilingualCheckbox');
   const languageDropdown = document.getElementById('languageDropdown');
   const taskDropdown = document.getElementById('taskDropdown');
   const modelSizeDropdown = document.getElementById('modelSizeDropdown');
@@ -32,14 +31,6 @@ 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("selectedLanguage", ({ selectedLanguage: storedLanguage }) => {
     if (storedLanguage !== undefined) {
       languageDropdown.value = storedLanguage;
@@ -73,7 +64,7 @@ document.addEventListener("DOMContentLoaded", function () {
 
     // Send a message to the background script to start capturing
     let host = "localhost";
-    let port = "5901";
+    let port = "9090";
     const useCollaboraServer = useServerCheckbox.checked;
     if (useCollaboraServer){
       host = "transcription.kurg.org"
@@ -86,7 +77,6 @@ document.addEventListener("DOMContentLoaded", function () {
         tabId: currentTab.id,
         host: host,
         port: port,
-        useMultilingual: useMultilingualCheckbox.checked,
         language: selectedLanguage,
         task: selectedTask,
         modelSize: selectedModelSize
@@ -129,9 +119,9 @@ document.addEventListener("DOMContentLoaded", function () {
     startButton.disabled = isCapturing;
     stopButton.disabled = !isCapturing;
     useServerCheckbox.disabled = isCapturing; 
-    useMultilingualCheckbox.disabled = isCapturing;
     modelSizeDropdown.disabled = isCapturing;
-    
+    languageDropdown.disabled = isCapturing;
+    taskDropdown.disabled = isCapturing; 
     startButton.classList.toggle("disabled", isCapturing);
     stopButton.classList.toggle("disabled", !isCapturing);
   }
@@ -142,18 +132,6 @@ 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 });
-  });
-
   languageDropdown.addEventListener('change', function() {
     if (languageDropdown.value === "") {
       selectedLanguage = null;

Audio-Transcription-Firefox/README.md

@@ -24,7 +24,6 @@ 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.

Audio-Transcription-Firefox/content.js

@@ -66,19 +66,15 @@ 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,
-            model_size: data.modelSize
+            model: data.modelSize
         })
       );
     };
@@ -201,7 +197,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

@@ -16,113 +16,109 @@
     <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>
-  </div>
   <div class="dropdown-container">
     <label for="languageDropdown">Select Language:</label>
-    <select id="languageDropdown" disabled>
+    <select id="languageDropdown">
-      <option value="">Select Language</option>
+      <option value="" selected>Automatically detect</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>
       <option value="af">Afrikaans</option>
-      <option value="oc">Occitan</option>
+      <option value="sq">Albanian</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="am">Amharic</option>
-      <option value="yi">Yiddish</option>
+      <option value="ar">Arabic</option>
-      <option value="lo">Lao</option>
+      <option value="hy">Armenian</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="as">Assamese</option>
-      <option value="tt">Tatar</option>
+      <option value="az">Azerbaijani</option>
-      <option value="haw">Hawaiian</option>
-      <option value="ln">Lingala</option>
-      <option value="ha">Hausa</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">
@@ -136,14 +132,18 @@
   <div class="dropdown-container">
     <label for="modelSizeDropdown">Select Model Size:</label>
     <select id="modelSizeDropdown">
-      <option value="">Select Task</option>
+      <option value="">Select model</option>
-      <option value="tiny">Tiny</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,7 +3,6 @@ document.addEventListener("DOMContentLoaded", function() {
   const stopButton = document.getElementById("stopCapture");
 
   const useServerCheckbox = document.getElementById("useServerCheckbox");
-  const useMultilingualCheckbox = document.getElementById('useMultilingualCheckbox');
   const languageDropdown = document.getElementById('languageDropdown');
   const taskDropdown = document.getElementById('taskDropdown');
   const modelSizeDropdown = document.getElementById('modelSizeDropdown');
@@ -35,14 +34,6 @@ 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("selectedLanguage", ({ selectedLanguage: storedLanguage }) => {
     if (storedLanguage !== undefined) {
       languageDropdown.value = storedLanguage;
@@ -66,7 +57,7 @@ document.addEventListener("DOMContentLoaded", function() {
 
   startButton.addEventListener("click", function() {
     let host = "localhost";
-    let port = "5901";
+    let port = "9090";
     const useCollaboraServer = useServerCheckbox.checked;
 
     if (useCollaboraServer){
@@ -83,7 +74,6 @@ document.addEventListener("DOMContentLoaded", function() {
             data: {
               host: host,
               port: port,
-              useMultilingual: useMultilingualCheckbox.checked,
               language: selectedLanguage,
               task: selectedTask,
               modelSize: selectedModelSize
@@ -125,9 +115,9 @@ document.addEventListener("DOMContentLoaded", function() {
     startButton.disabled = isCapturing;
     stopButton.disabled = !isCapturing;
     useServerCheckbox.disabled = isCapturing;
-    useMultilingualCheckbox.disabled = isCapturing;
     modelSizeDropdown.disabled = isCapturing;
-
+    languageDropdown.disabled = isCapturing;
+    taskDropdown.disabled = isCapturing; 
     startButton.classList.toggle("disabled", isCapturing);
     stopButton.classList.toggle("disabled", !isCapturing);
   }
@@ -138,18 +128,6 @@ 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 });
-  });
-
   languageDropdown.addEventListener('change', function() {
     if (languageDropdown.value === "") {
       selectedLanguage = null;

Audio-Transcription-Firefox/style.css

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

README.md

@@ -8,7 +8,7 @@ Unlike traditional speech recognition systems that rely on continuous audio stre
 ## Installation
 - Install PyAudio and ffmpeg
 ```bash
- bash setup.sh
+ bash scripts/setup.sh
 ```
 
 - Install whisper-live from pip
@@ -16,61 +16,79 @@ 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
+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)
-```python
+
- from whisper_live.server import TranscriptionServer
+### Running the Server
- server = TranscriptionServer()
+- [Faster Whisper](https://github.com/SYSTRAN/faster-whisper) backend
- server.run("0.0.0.0", 9090)
+```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
+- 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.
-    - To transcribe an audio file:
+```bash
-    ```python
+# Run English only model
-      from whisper_live.client import TranscriptionClient
+python3 run_server.py -p 9090 \
-      client = TranscriptionClient(
+                      -b tensorrt \
-        "localhost",
+                      -trt /home/TensorRT-LLM/examples/whisper/whisper_small_en
-        9090,
-        is_multilingual=False,
-        lang="en",
-        translate=False,
-        model_size="small"
-      )
 
-      client("tests/jfk.wav")
+# Run Multilingual model
-    ```
+python3 run_server.py -p 9090 \
-    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 can also enable the multilingual feature, allowing transcription in multiple languages. The language option specifies the target language for 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.
+                      -b tensorrt \
+                      -trt /home/TensorRT-LLM/examples/whisper/whisper_small \
+                      -m
+```
 
-    - To transcribe from microphone:
-    ```python
-      from whisper_live.client import TranscriptionClient
-      client = TranscriptionClient(
-        "localhost",
-        9090,
-        is_multilingual=True,
-        lang="hi",
-        translate=True,
-        model_size="small"
-      )
-      client()
-    ```
-    This command captures audio from the microphone and sends it to the server for transcription. It uses the multilingual option with `hi` as the selected language, enabling the multilingual feature and specifying the target language and task. We use whisper `small` by default but can be changed to any other option based on the requirements and the hardware running the server.
 
-    - To transcribe from a HLS stream:
+### Running the Client
-    ```python
+- To transcribe an audio file:
-      client = TranscriptionClient(host, port, is_multilingual=True, lang="en", translate=False) 
+```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") 
+from whisper_live.client import TranscriptionClient
-    ```
+client = TranscriptionClient(
-    This command streams audio into the server from a HLS stream. It uses the same options as the previous command, enabling the multilingual feature and specifying the target language and task.
+  "localhost",
+  9090,
+  lang="en",
+  translate=False,
+  model="small"
+)
+
+client("tests/jfk.wav")
+```
+This command transcribes the specified audio file (audio.wav) using the Whisper model. 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.
+
+- To transcribe from microphone:
+```python
+from whisper_live.client import TranscriptionClient
+client = TranscriptionClient(
+  "localhost",
+  9090,
+  lang="hi",
+  translate=True,
+  model="small"
+)
+client()
+```
+This command captures audio from the microphone and sends it to the server for transcription. It uses the multilingual model with `hi` as the selected language. We use whisper `small` by default but can be changed to any other option based on the requirements and the hardware running the server.
+
+- To transcribe from a HLS stream:
+```python
+from whisper_live.client import TranscriptionClient
+client = TranscriptionClient(host, port, lang="en", translate=False) 
+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") 
+```
+This command streams audio into the server from a HLS stream. It uses the same options as the previous command, using the multilingual model and specifying the target language and task.
 
 ## Transcribe audio from browser
-- Run the server
+- Run the server with your desired backend as shown [here](https://github.com/collabora/WhisperLive?tab=readme-ov-file#running-the-server)
-```python
- from whisper_live.server import TranscriptionServer
- server = TranscriptionServer()
- server.run("0.0.0.0", 9090)
-```
-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.
@@ -80,21 +98,24 @@ This would start the websocket server on port ```9090```.
 
 ## Whisper Live Server in Docker
 - GPU
-```bash
+  - Faster-Whisper
- docker build . -t whisper-live -f docker/Dockerfile.gpu
+  ```bash
- docker run -it --gpus all -p 9090:9090 whisper-live:latest
+  docker build . -t whisper-live -f docker/Dockerfile.gpu
-```
+  docker run -it --gpus all -p 9090:9090 whisper-live:latest
+  ```
+
+  - TensorRT. Follow [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md) in order to setup docker and use TensorRT backend. We provide a pre-built docker image which has TensorRT-LLM built and ready to use.
 
 - CPU
 ```bash
- docker build . -t whisper-live -f docker/Dockerfile.cpu
+docker build . -t whisper-live -f docker/Dockerfile.cpu
- docker run -it -p 9090:9090 whisper-live:latest
+docker run -it -p 9090:9090 whisper-live: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
 

TensorRT_whisper.md

@@ -0,0 +1,67 @@
+# Whisper-TensorRT
+We have only tested the TensorRT backend in docker so, we recommend docker for a smooth TensorRT backend setup.
+**Note**: We use [our fork to setup TensorRT](https://github.com/makaveli10/TensorRT-LLM)
+
+## Installation
+- Install [docker](https://docs.docker.com/engine/install/)
+- Install [nvidia-container-toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
+
+- Clone this repo.
+```bash
+git clone https://github.com/collabora/WhisperLive.git
+cd WhisperLive
+```
+
+- Pull the TensorRT-LLM docker image which we prebuilt for WhisperLive TensorRT backend.
+```bash
+docker pull ghcr.io/collabora/whisperbot-base:latest
+```
+
+- Next, we run the docker image and mount WhisperLive repo to the containers `/home` directory.
+```bash
+docker run -it --gpus all --shm-size=8g \
+       --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \
+       -v /path/to/WhisperLive:/home/WhisperLive \
+       ghcr.io/collabora/whisperbot-base:latest
+```
+
+- Make sure to test the installation. 
+```bash
+# export ENV=${ENV:-/etc/shinit_v2} 
+# source $ENV
+python -c "import torch; import tensorrt; import tensorrt_llm"
+```
+**NOTE**: Uncomment and update library paths if imports fail.
+
+## Whisper TensorRT Engine
+- We build `small.en` and `small` multilingual TensorRT engine. The script logs the path of the directory with Whisper TensorRT engine. We need the model_path to run the server.
+```bash
+# convert small.en
+bash scripts/build_whisper_tensorrt.sh /root/TensorRT-LLM-examples small.en
+
+# convert small multilingual model
+bash scripts/build_whisper_tensorrt.sh /root/TensorRT-LLM-examples small
+```
+
+## Run WhisperLive Server with TensorRT Backend
+```bash
+cd /home/WhisperLive
+
+# Install requirements
+bash scripts/setup.sh
+pip install -r requirements/server.txt
+
+# Required to create mel spectogram
+wget --directory-prefix=assets assets/mel_filters.npz https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz
+
+# Run English only model
+python3 run_server.py --port 9090 \
+                      --backend tensorrt \
+                      --trt_model_path "path/to/whisper_trt/from/build/step"
+
+# Run Multilingual model
+python3 run_server.py --port 9090 \
+                      --backend tensorrt \
+                      --trt_model_path "path/to/whisper_trt/from/build/step" \
+                      --trt_multilingual
+```

docker/Dockerfile.cpu

@@ -33,7 +33,7 @@ RUN apt install python3-pip -y
 RUN mkdir /app
 WORKDIR /app
 
-COPY setup.sh /app
+COPY scripts/setup.sh /app
 COPY requirements/ /app
 
 RUN bash setup.sh

docker/Dockerfile.gpu

@@ -33,7 +33,7 @@ RUN apt install python3-pip -y
 RUN mkdir /app
 WORKDIR /app
 
-COPY setup.sh /app
+COPY scripts/setup.sh /app
 COPY requirements/ /app
 
 RUN apt update --fix-missing

requirements/server.txt

@@ -1,7 +1,10 @@
-PyAudio
 faster-whisper==0.10.0
---extra-index-url https://download.pytorch.org/whl/cu111
+torch
-torch==1.10.1
-torchaudio==0.10.1
 websockets
-onnxruntime==1.16.0
+onnxruntime==1.16.0
+numba
+openai-whisper
+kaldialign
+soundfile
+ffmpeg-python
+scipy

run_server.py

@@ -1,5 +1,38 @@
+import argparse
 from whisper_live.server import TranscriptionServer
 
 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.')
+    args = parser.parse_args()
+
+    if args.backend == "tensorrt":
+        if args.trt_model_path is None:
+            raise ValueError("Please Provide a valid tensorrt model path")
+
     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
+    )

scripts/build_whisper_tensorrt.sh

@@ -0,0 +1,77 @@
+#!/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"
+            ;;
+        *)
+            echo "Invalid model name: $model_name"
+            exit 1
+            ;;
+    esac
+
+    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 output_dir="whisper_${model_name//./_}"
+    echo "$output_dir"
+    echo "Running build script for $model_name with output directory $output_dir"
+    python3 build.py --output_dir "$output_dir" --use_gpt_attention_plugin --use_gemm_plugin --use_bert_attention_plugin --model_name "$model_name"
+    echo "Whisper $model_name TensorRT engine built."
+    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}"
+
+cd $1/whisper
+pip install --no-deps -r requirements.txt
+
+download_and_build_model "$model_name"

setup.py

@@ -45,10 +45,14 @@ setup(name="whisper-live",
         "torch",
         "torchaudio",
         "websockets",
-        "onnxruntime",
+        "onnxruntime==1.16.0",
         "ffmpeg-python",
         "scipy",
         "websocket-client",
+        "numba",
+        "openai-whisper",
+        "kaldialign",
+        "soundfile",
       ],
       python_requires=">=3.8"
 )

whisper_live/__version__.py

@@ -1 +1 @@
-__version__="0.0.11"
+__version__="0.1.0"

whisper_live/client.py

@@ -13,6 +13,29 @@ import uuid
 import time
 
 
+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
@@ -50,7 +73,13 @@ class Client:
     INSTANCES = {}
 
     def __init__(
-        self, host=None, port=None, is_multilingual=False, lang=None, translate=False, model_size="small"
+        self,
+        host=None,
+        port=None,
+        lang=None,
+        translate=False,
+        model="small",
+        srt_file_path="output.srt"
     ):
         """
         Initializes a Client instance for audio recording and streaming to a server.
@@ -62,27 +91,25 @@ 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. Default is None.
-            lang (str, optional): The selected language for transcription when multilingual is disabled. Default is None.
             translate (bool, optional): Specifies if the task is translation. Default is False.
         """
-        self.chunk = 1024
+        self.chunk = 4096
         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.disconnect_if_no_response_for = 15
-        self.multilingual = is_multilingual
         self.language = lang
-        self.model_size = model_size
+        self.model = model
         self.server_error = False
+        self.srt_file_path = srt_file_path
+
         if translate:
             self.task = "translate"
 
@@ -120,6 +147,7 @@ class Client:
         self.ws_thread.start()
 
         self.frames = b""
+        self.transcript = []
         print("[INFO]: * recording")
 
     def on_message(self, ws, message):
@@ -159,6 +187,8 @@ class Client:
 
         if "message" in message.keys() and message["message"] == "SERVER_READY":
             self.recording = True
+            self.server_backend = message["backend"]
+            print(f"[INFO]: Server Running with backend {self.server_backend}")
             return
 
         if "language" in message.keys():
@@ -174,12 +204,21 @@ class Client:
 
         message = message["segments"]
         text = []
-        if len(message):
+        n_segments = len(message)
-            for seg in message:
+
+        if n_segments:
+            for i, seg in enumerate(message):
                 if text and text[-1] == seg["text"]:
                     # already got it
                     continue
                 text.append(seg["text"])
+
+                if i == n_segments-1: 
+                    self.last_segment = seg
+                elif self.server_backend == "faster_whisper":
+                    if not len(self.transcript) or float(seg['start']) >= float(self.transcript[-1]['end']):
+                        self.transcript.append(seg)
+
         # keep only last 3
         if len(text) > 3:
             text = text[-3:]
@@ -203,24 +242,21 @@ class Client:
         """
         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_size": self.model_size,
+                    "model": self.model,
                 }
             )
         )
@@ -295,6 +331,9 @@ class Client:
                 assert self.last_response_recieved
                 while time.time() - self.last_response_recieved < self.disconnect_if_no_response_for:
                     continue
+
+                if self.server_backend == "faster_whisper":
+                    self.write_srt_file(self.srt_file_path)
                 self.stream.close()
                 self.close_websocket()
 
@@ -304,6 +343,8 @@ class Client:
                 self.stream.close()
                 self.p.terminate()
                 self.close_websocket()
+                if self.server_backend == "faster_whisper":
+                    self.write_srt_file(self.srt_file_path)
                 print("[INFO]: Keyboard interrupt.")
 
     def close_websocket(self):
@@ -431,6 +472,8 @@ class Client:
                     t.start()
                     n_audio_file += 1
                     self.frames = b""
+            if self.server_backend == "faster_whisper":
+                self.write_srt_file(self.srt_file_path)
 
         except KeyboardInterrupt:
             if len(self.frames):
@@ -444,6 +487,8 @@ class Client:
             self.close_websocket()
 
             self.write_output_recording(n_audio_file, out_file)
+            if self.server_backend == "faster_whisper":
+                self.write_srt_file(self.srt_file_path)
 
     def write_output_recording(self, n_audio_file, out_file):
         """
@@ -480,6 +525,10 @@ class Client:
                 os.remove(in_file)
         wavfile.close()
 
+    def write_srt_file(self, output_path="output.srt"):
+        self.transcript.append(self.last_segment)
+        create_srt_file(self.transcript, output_path)
+
 
 class TranscriptionClient:
     """
@@ -491,8 +540,7 @@ 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. Default is None, which defaults to English ('en').
-        lang (str, optional): The primary language for transcription (used if `is_multilingual` is False). Default is None, which defaults to English ('en').
         translate (bool, optional): Indicates whether translation tasks are required (default is False).
 
     Attributes:
@@ -501,12 +549,18 @@ 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, model_size="small"):
+    def __init__(self,
-        self.client = Client(host, port, is_multilingual, lang, translate, model_size)
+        host,
+        port,
+        lang=None,
+        translate=False,
+        model="small",
+    ):
+        self.client = Client(host, port, lang, translate, model)
 
     def __call__(self, audio=None, hls_url=None):
         """

whisper_live/server.py

@@ -1,3 +1,4 @@
+import os
 import websockets
 import time
 import threading
@@ -5,14 +6,24 @@ import json
 import textwrap
 
 import logging
-# logging.basicConfig(level = logging.INFO)
+logging.basicConfig(level = logging.INFO)
 
 from websockets.sync.server import serve
 
 import torch
 import numpy as np
-import time
+import queue
+
+from whisper_live.vad import VoiceActivityDetection
+from scipy.io.wavfile import write
+import functools
+
+from whisper_live.vad import VoiceActivityDetection
 from whisper_live.transcriber import WhisperModel
+try:
+    from whisper_live.transcriber_tensorrt import WhisperTRTLLM
+except Exception as e:
+    logging.warn("cannot import WhisperTRTLLM")
 
 
 class TranscriptionServer:
@@ -34,7 +45,7 @@ class TranscriptionServer:
 
     def __init__(self):
         # voice activity detection model
-
+        
         self.clients = {}
         self.websockets = {}
         self.clients_start_time = {}
@@ -58,7 +69,12 @@ class TranscriptionServer:
 
         return wait_time / 60
 
-    def recv_audio(self, websocket):
+    def recv_audio(self,
+                   websocket,
+                   backend="faster_whisper",
+                   faster_whisper_custom_model_path=None,
+                   whisper_tensorrt_path=None,
+                   trt_multilingual=False):
         """
         Receive audio chunks from a client in an infinite loop.
         
@@ -75,10 +91,19 @@ class TranscriptionServer:
 
         Args:
             websocket (WebSocket): The WebSocket connection for the client.
-        
+            backend (str): The backend to run the server with.
+            faster_whisper_custom_model_path (str): path to custom faster whisper model.
+            whisper_tensorrt_path (str): Required for tensorrt backend.
+            trt_multilingual(bool): Only used for tensorrt, True if multilingual model.
+
         Raises:
             Exception: If there is an error during the audio frame processing.
         """
+        self.backend = backend
+        if self.backend == "tensorrt":
+            self.vad_model = VoiceActivityDetection()
+            self.vad_threshold = 0.5
+
         logging.info("New client connected")
         options = websocket.recv()
         options = json.loads(options)
@@ -96,25 +121,76 @@ class TranscriptionServer:
             del websocket
             return
 
-        client = ServeClient(
+        if self.backend == "tensorrt":
-            websocket,
+            try:
-            multilingual=options["multilingual"],
+                import tensorrt as trt
-            language=options["language"],
+                import tensorrt_llm
-            task=options["task"],
+                self.backend = "tensorrt"
-            client_uid=options["uid"],
+                client = ServeClientTensorRT(
-            model_size=options["model_size"],
+                    websocket,
-            initial_prompt=options.get("initial_prompt"),
+                    multilingual=trt_multilingual,
-            vad_parameters=options.get("vad_parameters")
+                    language=options["language"],
-        )
+                    task=options["task"],
+                    client_uid=options["uid"],
+                    model=whisper_tensorrt_path
+                )
+                logging.info(f"Running TensorRT backend.")
+            except Exception as e:
+                self.client_uid = options["uid"]
+                websocket.send(
+                    json.dumps(
+                        {
+                            "uid": self.client_uid,
+                            "status": "ERROR",
+                            "message": f"TensorRT-LLM not supported on Server yet. Reverting to available backend: 'faster_whisper'"
+                        }
+                    )
+                )
+                self.backend = "faster_whisper"
+
+        if self.backend == "faster_whisper":
+            # validate custom model
+            if faster_whisper_custom_model_path is not None and os.path.exists(faster_whisper_custom_model_path):
+                logging.info(f"Using custom model {faster_whisper_custom_model_path}")
+                options["model"] = faster_whisper_custom_model_path
+            client = ServeClientFasterWhisper(
+                websocket,
+                language=options["language"],
+                task=options["task"],
+                client_uid=options["uid"],
+                model=options["model"],
+                initial_prompt=options.get("initial_prompt"),
+                vad_parameters=options.get("vad_parameters")
+            )
+            logging.info(f"Running faster_whisper backend.")
         
         self.clients[websocket] = client
         self.clients_start_time[websocket] = time.time()
+        no_voice_activity_chunks = 0
 
         while True:
             try:
                 frame_data = websocket.recv()
                 frame_np = np.frombuffer(frame_data, dtype=np.float32)
 
+                # VAD, for faster_whisper VAD model is already integrated
+                if self.backend == "tensorrt":
+                    try:
+                        speech_prob = self.vad_model(torch.from_numpy(frame_np.copy()), self.RATE).item()
+                        if speech_prob < self.vad_threshold:
+                            no_voice_activity_chunks += 1
+                            if no_voice_activity_chunks > 3:
+                                if not self.clients[websocket].eos:
+                                    self.clients[websocket].set_eos(True)
+                                time.sleep(0.1)    # Sleep 100m; wait some voice activity.
+                            continue
+                        no_voice_activity_chunks = 0
+                        self.clients[websocket].set_eos(False)
+
+                    except Exception as e:
+                        logging.error(e)
+                        return
+                
                 self.clients[websocket].add_frames(frame_np)
 
                 elapsed_time = time.time() - self.clients_start_time[websocket]
@@ -129,15 +205,21 @@ class TranscriptionServer:
                     break
 
             except Exception as e:
-                logging.info(f"[ERROR]: Client with uid '{self.clients[websocket].client_uid}' Disconnected.")
+                logging.error(e)
-                if self.clients[websocket].model_size is not None:
+                self.clients[websocket].cleanup()
-                    self.clients[websocket].cleanup()
                 self.clients.pop(websocket)
                 self.clients_start_time.pop(websocket)
                 del websocket
                 break
 
-    def run(self, host, port=9090):
+    def run(self, 
+            host, 
+            port=9090, 
+            backend="tensorrt", 
+            faster_whisper_custom_model_path=None,
+            whisper_tensorrt_path=None, 
+            trt_multilingual=False
+        ):
         """
         Run the transcription server.
 
@@ -145,11 +227,111 @@ class TranscriptionServer:
             host (str): The host address to bind the server.
             port (int): The port number to bind the server.
         """
-        with serve(self.recv_audio, host, port) as server:
+        with serve(
+            functools.partial(
+                self.recv_audio,
+                backend=backend,
+                faster_whisper_custom_model_path=faster_whisper_custom_model_path,
+                whisper_tensorrt_path=whisper_tensorrt_path,
+                trt_multilingual=trt_multilingual
+            ),
+            host,
+            port
+        ) as server:
             server.serve_forever()
 
 
-class ServeClient:
+class ServeClientBase(object):
+    RATE = 16000
+    SERVER_READY = "SERVER_READY"
+    DISCONNECT = "DISCONNECT"
+
+    def __init__(self, client_uid, websocket):
+        self.client_uid = client_uid
+        self.websocket = websocket
+        self.data = b""
+        self.frames = b""
+        self.timestamp_offset = 0.0
+        self.frames_np = None
+        self.frames_offset = 0.0
+        self.text = []
+        self.current_out = ''
+        self.prev_out = ''
+        self.t_start=None
+        self.exit = False
+        self.same_output_threshold = 0
+        self.show_prev_out_thresh = 5   # if pause(no output from whisper) show previous output for 5 seconds
+        self.add_pause_thresh = 3       # add a blank to segment list as a pause(no speech) for 3 seconds
+        self.transcript = []
+        self.send_last_n_segments = 10
+
+        # text formatting
+        self.wrapper = textwrap.TextWrapper(width=50)
+        self.pick_previous_segments = 2
+
+        # threading
+        self.lock = threading.Lock()
+    
+    def add_frames(self, frame_np):
+        """
+        Add audio frames to the ongoing audio stream buffer.
+
+        This method is responsible for maintaining the audio stream buffer, allowing the continuous addition
+        of audio frames as they are received. It also ensures that the buffer does not exceed a specified size
+        to prevent excessive memory usage.
+
+        If the buffer size exceeds a threshold (45 seconds of audio data), it discards the oldest 30 seconds
+        of audio data to maintain a reasonable buffer size. If the buffer is empty, it initializes it with the provided
+        audio frame. The audio stream buffer is used for real-time processing of audio data for transcription.
+
+        Args:
+            frame_np (numpy.ndarray): The audio frame data as a NumPy array.
+
+        """
+        self.lock.acquire()
+        if self.frames_np is not None and self.frames_np.shape[0] > 45*self.RATE:
+            self.frames_offset += 30.0
+            self.frames_np = self.frames_np[int(30*self.RATE):]
+        if self.frames_np is None:
+            self.frames_np = frame_np.copy()
+        else:
+            self.frames_np = np.concatenate((self.frames_np, frame_np), axis=0)
+        self.lock.release()
+
+    def speech_to_text(self):
+        raise NotImplementedError("Please implement in child Class.")
+    
+    def disconnect(self):
+        """
+        Notify the client of disconnection and send a disconnect message.
+
+        This method sends a disconnect message to the client via the WebSocket connection to notify them
+        that the transcription service is disconnecting gracefully.
+
+        """
+        self.websocket.send(
+            json.dumps(
+                {
+                    "uid": self.client_uid,
+                    "message": self.DISCONNECT
+                }
+            )
+        )
+    
+    def cleanup(self):
+        """
+        Perform cleanup tasks before exiting the transcription service.
+
+        This method performs necessary cleanup tasks, including stopping the transcription thread, marking
+        the exit flag to indicate the transcription thread should exit gracefully, and destroying resources
+        associated with the transcription process.
+
+        """
+        logging.info("Cleaning up.")
+        self.exit = True
+
+
+class ServeClientTensorRT(ServeClientBase):
     """
     Attributes:
         RATE (int): The audio sampling rate (constant) set to 16000.
@@ -178,21 +360,15 @@ class ServeClient:
         pick_previous_segments (int): Number of previous segments to include in the output.
         websocket: The WebSocket connection for the client.
     """
-    RATE = 16000
-    SERVER_READY = "SERVER_READY"
-    DISCONNECT = "DISCONNECT"
-
     def __init__(
         self,
         websocket,
         task="transcribe",
         device=None,
         multilingual=False,
-        language=None,
+        language=None, 
         client_uid=None,
-        model_size="small",
+        model=None
-        initial_prompt=None,
-        vad_parameters=None
         ):
         """
         Initialize a ServeClient instance.
@@ -209,86 +385,44 @@ class ServeClient:
             client_uid (str, optional): A unique identifier for the client. Defaults to None.
 
         """
-        self.client_uid = client_uid
+        super().__init__(client_uid, websocket)
-        self.data = b""
+        self.language = language if multilingual else "en"
-        self.frames = b""
-        self.model_sizes = [
-            "tiny", "base", "small", "medium", "large-v2", "large-v3"
-        ]
-        self.multilingual = multilingual
-        self.model_size = self.get_model_size(model_size)
-        self.language = language if self.multilingual else "en"
         self.task = task
-        self.websocket = websocket
+        self.eos = False
-        self.initial_prompt = initial_prompt
+        self.transcriber = WhisperTRTLLM(
-        self.vad_parameters = vad_parameters or {"threshold": 0.5}
+            model, 
-        
+            assets_dir="assets", 
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+            device="cuda",
-        
+            is_multilingual=multilingual,
-        if self.model_size == None:
+            language=self.language,
-            return
+            task=self.task
-        
-        self.transcriber = WhisperModel(
-            self.model_size, 
-            device=device,
-            compute_type="int8" if device=="cpu" else "float16", 
-            local_files_only=False,
         )
-        
+        self.warmup()
-        self.timestamp_offset = 0.0
-        self.frames_np = None
-        self.frames_offset = 0.0
-        self.text = []
-        self.current_out = ''
-        self.prev_out = ''
-        self.t_start=None
-        self.exit = False
-        self.same_output_threshold = 0
-        self.show_prev_out_thresh = 5   # if pause(no output from whisper) show previous output for 5 seconds
-        self.add_pause_thresh = 3       # add a blank to segment list as a pause(no speech) for 3 seconds
-        self.transcript = []
-        self.send_last_n_segments = 10
-
-        # text formatting
-        self.wrapper = textwrap.TextWrapper(width=50)
-        self.pick_previous_segments = 2
 
         # threading
         self.trans_thread = threading.Thread(target=self.speech_to_text)
         self.trans_thread.start()
+
         self.websocket.send(
             json.dumps(
                 {
                     "uid": self.client_uid,
-                    "message": self.SERVER_READY
+                    "message": self.SERVER_READY,
+                    "backend": "tensorrt"
                 }
             )
         )
+
+    def warmup(self, warmup_steps=10):
+        logging.info("[INFO:] Warming up TensorRT engine..")
+        mel, _ = self.transcriber.log_mel_spectrogram("tests/jfk.flac")
+        for i in range(warmup_steps):
+            self.transcriber.transcribe(mel)
     
-    def get_model_size(self, model_size):
+    def set_eos(self, eos):
-        """
+        self.lock.acquire()
-        Returns the whisper model size based on multilingual.
+        self.eos = eos
-        """
+        self.lock.release()
-        if model_size not in self.model_sizes:
-            self.websocket.send(
-                json.dumps(
-                    {
-                        "uid": self.client_uid,
-                        "status": "ERROR",
-                        "message": f"Invalid model size {model_size}. Available choices: {self.model_sizes}"
-                    }
-                )
-            )
-            return None
-        
-        if model_size in ["large-v2", "large-v3"]:
-            self.multilingual = True
-            return model_size
-
-        if not self.multilingual:
-            model_size = model_size + ".en"
-
-        return model_size
     
     def add_frames(self, frame_np):
         """
@@ -306,6 +440,7 @@ class ServeClient:
             frame_np (numpy.ndarray): The audio frame data as a NumPy array.
 
         """
+        self.lock.acquire()
         if self.frames_np is not None and self.frames_np.shape[0] > 45*self.RATE:
             self.frames_offset += 30.0
             self.frames_np = self.frames_np[int(30*self.RATE):]
@@ -313,7 +448,199 @@ class ServeClient:
             self.frames_np = frame_np.copy()
         else:
             self.frames_np = np.concatenate((self.frames_np, frame_np), axis=0)
+        self.lock.release()
 
+    def speech_to_text(self):
+        """
+        Process an audio stream in an infinite loop, continuously transcribing the speech.
+
+        This method continuously receives audio frames, performs real-time transcription, and sends
+        transcribed segments to the client via a WebSocket connection.
+
+        If the client's language is not detected, it waits for 30 seconds of audio input to make a language prediction.
+        It utilizes the Whisper ASR model to transcribe the audio, continuously processing and streaming results. Segments
+        are sent to the client in real-time, and a history of segments is maintained to provide context.Pauses in speech 
+        (no output from Whisper) are handled by showing the previous output for a set duration. A blank segment is added if 
+        there is no speech for a specified duration to indicate a pause.
+
+        Raises:
+            Exception: If there is an issue with audio processing or WebSocket communication.
+
+        """
+        while True:
+            if self.exit:
+                logging.info("Exiting speech to text thread")
+                break
+            
+            if self.frames_np is None:
+                time.sleep(0.02)    # wait for any audio to arrive
+                continue
+
+            # clip audio if the current chunk exceeds 30 seconds, this basically implies that
+            # no valid segment for the last 30 seconds from whisper
+            if self.frames_np[int((self.timestamp_offset - self.frames_offset)*self.RATE):].shape[0] > 25 * self.RATE:
+                duration = self.frames_np.shape[0] / self.RATE
+                self.timestamp_offset = self.frames_offset + duration - 5
+    
+            samples_take = max(0, (self.timestamp_offset - self.frames_offset)*self.RATE)
+            input_bytes = self.frames_np[int(samples_take):].copy()
+            duration = input_bytes.shape[0] / self.RATE
+            if duration<0.4:
+                continue
+
+            try:
+                input_sample = input_bytes.copy()
+                logging.info(f"[WhisperTensorRT:] Processing audio with duration: {duration}")
+                mel, duration = self.transcriber.log_mel_spectrogram(input_sample)
+                last_segment = self.transcriber.transcribe(mel)
+                segments = []
+                if len(last_segment):
+                    if len(self.transcript) < self.send_last_n_segments:
+                        segments = self.transcript[:].copy()
+                    else:
+                        segments = self.transcript[-self.send_last_n_segments:].copy()
+                    if last_segment is not None:
+                        segments.append({"text": last_segment})
+                    try:
+                        self.websocket.send(
+                            json.dumps({
+                                "uid": self.client_uid,
+                                "segments": segments,
+                            })
+                        )
+
+                        if self.eos:
+                            if not len(self.transcript):
+                                self.transcript.append({"text": last_segment + " "})
+                            elif self.transcript[-1]["text"].strip() != last_segment:
+                                self.transcript.append({"text": last_segment + " "})
+                            self.timestamp_offset += duration
+
+                            
+                    except Exception as e:
+                        logging.error(f"[ERROR]: {e}")
+
+            except Exception as e:
+                logging.error(f"[ERROR]: {e}")
+
+
+class ServeClientFasterWhisper(ServeClientBase):
+    """
+    Attributes:
+        RATE (int): The audio sampling rate (constant) set to 16000.
+        SERVER_READY (str): A constant message indicating that the server is ready.
+        DISCONNECT (str): A constant message indicating that the client should disconnect.
+        client_uid (str): A unique identifier for the client.
+        data (bytes): Accumulated audio data.
+        frames (bytes): Accumulated audio frames.
+        language (str): The language for transcription.
+        task (str): The task type, e.g., "transcribe."
+        transcriber (WhisperModel): The Whisper model for speech-to-text.
+        timestamp_offset (float): The offset in audio timestamps.
+        frames_np (numpy.ndarray): NumPy array to store audio frames.
+        frames_offset (float): The offset in audio frames.
+        text (list): List of transcribed text segments.
+        current_out (str): The current incomplete transcription.
+        prev_out (str): The previous incomplete transcription.
+        t_start (float): Timestamp for the start of transcription.
+        exit (bool): A flag to exit the transcription thread.
+        same_output_threshold (int): Threshold for consecutive same output segments.
+        show_prev_out_thresh (int): Threshold for showing previous output segments.
+        add_pause_thresh (int): Threshold for adding a pause (blank) segment.
+        transcript (list): List of transcribed segments.
+        send_last_n_segments (int): Number of last segments to send to the client.
+        wrapper (textwrap.TextWrapper): Text wrapper for formatting text.
+        pick_previous_segments (int): Number of previous segments to include in the output.
+        websocket: The WebSocket connection for the client.
+    """
+    def __init__(
+        self,
+        websocket,
+        task="transcribe",
+        device=None,
+        language=None,
+        client_uid=None,
+        model="small.en",
+        initial_prompt=None,
+        vad_parameters=None,
+        ):
+        """
+        Initialize a ServeClient instance.
+        The Whisper model is initialized based on the client's language and device availability.
+        The transcription thread is started upon initialization. A "SERVER_READY" message is sent
+        to the client to indicate that the server is ready.
+
+        Args:
+            websocket (WebSocket): The WebSocket connection for the client.
+            task (str, optional): The task type, e.g., "transcribe." Defaults to "transcribe".
+            device (str, optional): The device type for Whisper, "cuda" or "cpu". Defaults to None.
+            language (str, optional): The language for transcription. Defaults to None.
+            client_uid (str, optional): A unique identifier for the client. Defaults to None.
+
+        """
+        super().__init__(client_uid, websocket)
+        self.model_sizes = [
+            "tiny", "tiny.en", "base", "base.en", "small", "small.en",
+            "medium", "medium.en", "large-v2", "large-v3",
+        ]
+        if not os.path.exists(model):
+            self.model_size_or_path = self.check_valid_model(model)
+        else:
+            self.model_size_or_path = model
+        self.language = "en" if self.model_size_or_path.endswith("en") else language
+        self.task = task
+        self.initial_prompt = initial_prompt
+        self.vad_parameters = vad_parameters or {"threshold": 0.5}
+        self.no_speech_thresh = 0.45
+        
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        if self.model_size_or_path == None:
+            return
+
+        self.transcriber = WhisperModel(
+            self.model_size_or_path, 
+            device=device,
+            compute_type="int8" if device=="cpu" else "float16", 
+            local_files_only=False,
+        )
+
+        # threading
+        self.trans_thread = threading.Thread(target=self.speech_to_text)
+        self.trans_thread.start()
+        self.websocket.send(
+            json.dumps(
+                {
+                    "uid": self.client_uid,
+                    "message": self.SERVER_READY,
+                    "backend": "faster_whisper"
+                }
+            )
+        )
+    
+    def check_valid_model(self, model_size):
+        """
+        Check if it's a valid whisper model size.
+
+        Args:
+            model_size (str): The name of the model size to check.
+
+        Returns:
+            str: The model size if valid, None otherwise.
+        """
+        if model_size not in self.model_sizes:
+            self.websocket.send(
+                json.dumps(
+                    {
+                        "uid": self.client_uid,
+                        "status": "ERROR",
+                        "message": f"Invalid model size {model_size}. Available choices: {self.model_sizes}"
+                    }
+                )
+            )
+            return None
+        return model_size
+    
     def speech_to_text(self):
         """
         Process an audio stream in an infinite loop, continuously transcribing the speech.
@@ -397,6 +724,7 @@ class ServeClient:
                         if time.time() - self.t_start > self.add_pause_thresh:
                             self.text.append('')
 
+                if not len(segments): continue
                 try:
                     self.websocket.send(
                         json.dumps({
@@ -449,6 +777,10 @@ class ServeClient:
                 text_ = s.text
                 self.text.append(text_)
                 start, end = self.timestamp_offset + s.start, self.timestamp_offset + min(duration, s.end)
+
+                if start >= end: continue
+                if s.no_speech_prob > self.no_speech_thresh: continue
+
                 self.transcript.append(self.format_segment(start, end, text_))
                 
                 offset = min(duration, s.end)
@@ -487,32 +819,3 @@ class ServeClient:
             self.timestamp_offset += offset
 
         return last_segment
-    
-    def disconnect(self):
-        """
-        Notify the client of disconnection and send a disconnect message.
-
-        This method sends a disconnect message to the client via the WebSocket connection to notify them
-        that the transcription service is disconnecting gracefully.
-
-        """
-        self.websocket.send(
-            json.dumps(
-                {
-                    "uid": self.client_uid,
-                    "message": self.DISCONNECT
-                }
-            )
-        )
-    
-    def cleanup(self):
-        """
-        Perform cleanup tasks before exiting the transcription service.
-
-        This method performs necessary cleanup tasks, including stopping the transcription thread, marking
-        the exit flag to indicate the transcription thread should exit gracefully, and destroying resources
-        associated with the transcription process.
-
-        """
-        logging.info("Cleaning up.")
-        self.exit = True

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

whisper_live/vad.py

@@ -0,0 +1,118 @@
+# original: https://github.com/snakers4/silero-vad/blob/master/utils_vad.py
+
+import os
+import subprocess
+import torch
+import numpy as np
+import onnxruntime
+
+
+class VoiceActivityDetection():
+
+    def __init__(self, force_onnx_cpu=True):
+        print("downloading ONNX model...")
+        path = self.download()
+        print("loading session")
+
+        opts = onnxruntime.SessionOptions()
+        opts.log_severity_level = 3
+
+        opts.inter_op_num_threads = 1
+        opts.intra_op_num_threads = 1
+
+        print("loading onnx model")
+        if force_onnx_cpu and 'CPUExecutionProvider' in onnxruntime.get_available_providers():
+            self.session = onnxruntime.InferenceSession(path, providers=['CPUExecutionProvider'], sess_options=opts)
+        else:
+            self.session = onnxruntime.InferenceSession(path, providers=['CUDAExecutionProvider'], sess_options=opts)
+
+        print("reset states")
+        self.reset_states()
+        self.sample_rates = [8000, 16000]
+
+    def _validate_input(self, x, sr: int):
+        if x.dim() == 1:
+            x = x.unsqueeze(0)
+        if x.dim() > 2:
+            raise ValueError(f"Too many dimensions for input audio chunk {x.dim()}")
+
+        if sr != 16000 and (sr % 16000 == 0):
+            step = sr // 16000
+            x = x[:,::step]
+            sr = 16000
+
+        if sr not in self.sample_rates:
+            raise ValueError(f"Supported sampling rates: {self.sample_rates} (or multiply of 16000)")
+
+        if sr / x.shape[1] > 31.25:
+            raise ValueError("Input audio chunk is too short")
+
+        return x, sr
+
+    def reset_states(self, batch_size=1):
+        self._h = np.zeros((2, batch_size, 64)).astype('float32')
+        self._c = np.zeros((2, batch_size, 64)).astype('float32')
+        self._last_sr = 0
+        self._last_batch_size = 0
+
+    def __call__(self, x, sr: int):
+
+        x, sr = self._validate_input(x, sr)
+        batch_size = x.shape[0]
+
+        if not self._last_batch_size:
+            self.reset_states(batch_size)
+        if (self._last_sr) and (self._last_sr != sr):
+            self.reset_states(batch_size)
+        if (self._last_batch_size) and (self._last_batch_size != batch_size):
+            self.reset_states(batch_size)
+
+        if sr in [8000, 16000]:
+            ort_inputs = {'input': x.numpy(), 'h': self._h, 'c': self._c, 'sr': np.array(sr, dtype='int64')}
+            ort_outs = self.session.run(None, ort_inputs)
+            out, self._h, self._c = ort_outs
+        else:
+            raise ValueError()
+
+        self._last_sr = sr
+        self._last_batch_size = batch_size
+
+        out = torch.tensor(out)
+        return out
+
+    def audio_forward(self, x, sr: int, num_samples: int = 512):
+        outs = []
+        x, sr = self._validate_input(x, sr)
+
+        if x.shape[1] % num_samples:
+            pad_num = num_samples - (x.shape[1] % num_samples)
+            x = torch.nn.functional.pad(x, (0, pad_num), 'constant', value=0.0)
+
+        self.reset_states(x.shape[0])
+        for i in range(0, x.shape[1], num_samples):
+            wavs_batch = x[:, i:i+num_samples]
+            out_chunk = self.__call__(wavs_batch, sr)
+            outs.append(out_chunk)
+
+        stacked = torch.cat(outs, dim=1)
+        return stacked.cpu()
+
+    @staticmethod
+    def download(model_url="https://github.com/snakers4/silero-vad/raw/master/files/silero_vad.onnx"):
+        target_dir = os.path.expanduser("~/.cache/whisper-live/")
+
+        # Ensure the target directory exists
+        os.makedirs(target_dir, exist_ok=True)
+
+        # Define the target file path
+        model_filename = os.path.join(target_dir, "silero_vad.onnx")
+
+        # 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