🔨 refactor whisper_live according to flake8

setup.py

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

tests/test_client.py

@@ -4,7 +4,8 @@ import scipy
 import websocket
 import unittest
 from unittest.mock import patch, MagicMock
-from whisper_live.client import TranscriptionClient, resample
+from whisper_live.client import TranscriptionClient
+from whisper_live.utils import resample
 
 
 class BaseTestCase(unittest.TestCase):

tests/test_server.py

@@ -14,26 +14,25 @@ from whisper.normalizers import EnglishTextNormalizer
 class TestTranscriptionServerInitialization(unittest.TestCase):
     def test_initialization(self):
         server = TranscriptionServer()
-        self.assertEqual(server.max_clients, 4)
+        self.assertEqual(server.client_manager.max_clients, 4)
-        self.assertEqual(server.max_connection_time, 600)
+        self.assertEqual(server.client_manager.max_connection_time, 600)
-        self.assertDictEqual(server.clients, {})
+        self.assertDictEqual(server.client_manager.clients, {})
-        self.assertDictEqual(server.websockets, {})
+        self.assertDictEqual(server.client_manager.start_times, {})
-        self.assertDictEqual(server.clients_start_time, {})
 
 
 class TestGetWaitTime(unittest.TestCase):
     def setUp(self):
         self.server = TranscriptionServer()
-        self.server.clients_start_time = {
+        self.server.client_manager.start_times = {
             'client1': time.time() - 120,
             'client2': time.time() - 300
         }
-        self.server.max_connection_time = 600
+        self.server.client_manager.max_connection_time = 600
 
     def test_get_wait_time(self):
-        expected_wait_time = (600 - (time.time() - self.server.clients_start_time['client2'])) / 60
+        expected_wait_time = (600 - (time.time() - self.server.client_manager.start_times['client2'])) / 60
-        print(self.server.get_wait_time(), expected_wait_time)
+        print(self.server.client_manager.get_wait_time(), expected_wait_time)
-        self.assertAlmostEqual(self.server.get_wait_time(), expected_wait_time, places=2)
+        self.assertAlmostEqual(self.server.client_manager.get_wait_time(), expected_wait_time, places=2)
 
 
 class TestServerConnection(unittest.TestCase):
@@ -50,7 +49,6 @@ class TestServerConnection(unittest.TestCase):
         })
         self.server.recv_audio(mock_websocket, "faster_whisper")
 
-    
     @mock.patch('websockets.WebSocketCommonProtocol')
     def test_recv_audio_exception_handling(self, mock_websocket):
         mock_websocket.recv.side_effect = [json.dumps({
@@ -63,7 +61,7 @@ class TestServerConnection(unittest.TestCase):
         with self.assertLogs(level="ERROR"):
             self.server.recv_audio(mock_websocket, "faster_whisper")
 
-        self.assertNotIn(mock_websocket, self.server.clients)
+        self.assertNotIn(mock_websocket, self.server.client_manager.clients)
 
 
 class TestServerInferenceAccuracy(unittest.TestCase):
@@ -84,7 +82,7 @@ class TestServerInferenceAccuracy(unittest.TestCase):
         self.mock_pyaudio.open.return_value = self.mock_stream
         self.metric = evaluate.load("wer")
         self.normalizer = EnglishTextNormalizer()
-        self.client  = TranscriptionClient(
+        self.client = TranscriptionClient(
             "localhost", "9090", model="base.en", lang="en",
         )
 
@@ -93,7 +91,7 @@ class TestServerInferenceAccuracy(unittest.TestCase):
         self.client("assets/jfk.flac")
         with open("output.srt", "r") as f:
             lines = f.readlines()
-            prediction = " ".join([l.strip() for l in lines[2::4]])
+            prediction = " ".join([line.strip() for line in lines[2::4]])
         prediction_normalized = self.normalizer(prediction)
         gt_normalized = self.normalizer(gt)
 

tests/test_vad.py

@@ -1,7 +1,6 @@
 import unittest
 import numpy as np
 import torch
-import scipy.io as sio
 from whisper_live.tensorrt_utils import load_audio
 from whisper_live.vad import VoiceActivityDetection
 

whisper_live/__version__.py

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

whisper_live/client.py

@@ -2,68 +2,14 @@ import os
 import wave
 
 import numpy as np
-import scipy
-import ffmpeg
 import pyaudio
 import threading
-import textwrap
 import json
 import websocket
 import uuid
 import time
-
+import ffmpeg
-
+import whisper_live.utils as utils
-def format_time(s):
-    """Convert seconds (float) to SRT time format."""
-    hours = int(s // 3600)
-    minutes = int((s % 3600) // 60)
-    seconds = int(s % 60)
-    milliseconds = int((s - int(s)) * 1000)
-    return f"{hours:02}:{minutes:02}:{seconds:02},{milliseconds:03}"
-
-def create_srt_file(segments, output_file):
-    with open(output_file, 'w', encoding='utf-8') as srt_file:
-        segment_number = 1
-        for segment in segments:
-            start_time = format_time(float(segment['start']))
-            end_time = format_time(float(segment['end']))
-            text = segment['text']
-            
-            srt_file.write(f"{segment_number}\n")
-            srt_file.write(f"{start_time} --> {end_time}\n")
-            srt_file.write(f"{text}\n\n")
-            
-            segment_number += 1
-
-
-def resample(file: str, sr: int = 16000):
-    """
-    # https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/audio.py#L22
-    Open an audio file and read as mono waveform, resampling as necessary,
-    save the resampled audio
-
-    Args:
-        file (str): The audio file to open
-        sr (int): The sample rate to resample the audio if necessary
-    
-    Returns:
-        resampled_file (str): The resampled audio file
-    """
-    try:
-        # This launches a subprocess to decode audio while down-mixing and resampling as necessary.
-        # Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
-        out, _ = (
-            ffmpeg.input(file, threads=0)
-            .output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=sr)
-            .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
-        )
-    except ffmpeg.Error as e:
-        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
-    np_buffer = np.frombuffer(out, dtype=np.int16)
-
-    resampled_file = f"{file.split('.')[0]}_resampled.wav"
-    scipy.io.wavfile.write(resampled_file, sr, np_buffer.astype(np.int16))
-    return resampled_file
 
 
 class Client:
@@ -150,6 +96,36 @@ class Client:
         self.transcript = []
         print("[INFO]: * recording")
 
+    def handle_status_messages(self, message_data):
+        """Handles server status messages."""
+        status = message_data["status"]
+        if status == "WAIT":
+            self.waiting = True
+            print(f"[INFO]: Server is full. Estimated wait time {round(message_data['message'])} minutes.")
+        elif status == "ERROR":
+            print(f"Message from Server: {message_data['message']}")
+            self.server_error = True
+        elif status == "WARNING":
+            print(f"Message from Server: {message_data['message']}")
+
+    def process_segments(self, segments):
+        """Processes transcript segments."""
+        text = []
+        for i, seg in enumerate(segments):
+            if not text or text[-1] != seg["text"]:
+                text.append(seg["text"])
+                if i == len(segments) - 1:
+                    self.last_segment = seg
+                elif (self.server_backend == "faster_whisper" and
+                      (not self.transcript or
+                        float(seg['start']) >= float(self.transcript[-1]['end']))):
+                    self.transcript.append(seg)
+
+        # Truncate to last 3 entries for brevity.
+        text = text[-3:]
+        # utils.clear_screen()
+        utils.print_transcript(text)
+
     def on_message(self, ws, message):
         """
         Callback function called when a message is received from the server.
@@ -171,14 +147,7 @@ class Client:
             return
 
         if "status" in message.keys():
-            if message["status"] == "WAIT":
+            self.handle_status_messages(message)
-                self.waiting = True
-                print(
-                    f"[INFO]:Server is full. Estimated wait time {round(message['message'])} minutes."
-                )
-            elif message["status"] == "ERROR":
-                print(f"Message from Server: {message['message']}")
-                self.server_error = True
             return
 
         if "message" in message.keys() and message["message"] == "DISCONNECT":
@@ -199,38 +168,8 @@ class Client:
             )
             return
 
-        if "segments" not in message.keys():
+        if "segments" in message.keys():
-            return
+            self.process_segments(message["segments"])
-
-        message = message["segments"]
-        text = []
-        n_segments = len(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:]
-        wrapper = textwrap.TextWrapper(width=60)
-        word_list = wrapper.wrap(text="".join(text))
-        # Print each line.
-        if os.name == "nt":
-            os.system("cls")
-        else:
-            os.system("clear")
-        for element in word_list:
-            print(element)
 
     def on_error(self, ws, error):
         print(f"[ERROR] WebSocket Error: {error}")
@@ -433,7 +372,6 @@ class Client:
 
         print("[INFO]: HLS stream processing finished.")
 
-
     def record(self, out_file="output_recording.wav"):
         """
         Record audio data from the input stream and save it to a WAV file.
@@ -448,7 +386,8 @@ class Client:
         the method combines all the saved audio chunks into the specified `out_file`.
 
         Args:
-            out_file (str, optional): The name of the output WAV file to save the entire recording. Default is "output_recording.wav".
+            out_file (str, optional): The name of the output WAV file to save the entire recording.
+                                      Default is "output_recording.wav".
 
         """
         n_audio_file = 0
@@ -458,7 +397,7 @@ class Client:
             for _ in range(0, int(self.rate / self.chunk * self.record_seconds)):
                 if not self.recording:
                     break
-                data = self.stream.read(self.chunk, exception_on_overflow = False)
+                data = self.stream.read(self.chunk, exception_on_overflow=False)
                 self.frames += data
 
                 audio_array = Client.bytes_to_float_array(data)
@@ -532,7 +471,7 @@ class Client:
 
     def write_srt_file(self, output_path="output.srt"):
         self.transcript.append(self.last_segment)
-        create_srt_file(self.transcript, output_path)
+        utils.create_srt_file(self.transcript, output_path)
 
 
 class TranscriptionClient:
@@ -558,13 +497,7 @@ class TranscriptionClient:
         transcription_client()
         ```
     """
-    def __init__(self,
+    def __init__(self, host, port, lang=None, translate=False, model="small"):
-        host,
-        port,
-        lang=None,
-        translate=False,
-        model="small",
-    ):
         self.client = Client(host, port, lang, translate, model)
 
     def __call__(self, audio=None, hls_url=None):
@@ -589,7 +522,7 @@ class TranscriptionClient:
         if hls_url is not None:
             self.client.process_hls_stream(hls_url)
         elif audio is not None:
-            resampled_file = resample(audio)
+            resampled_file = utils.resample(audio)
             self.client.play_file(resampled_file)
         else:
             self.client.record()

whisper_live/server.py

@@ -3,25 +3,80 @@ import time
 import threading
 import json
 import textwrap
-
+import functools
 import logging
-logging.basicConfig(level = logging.INFO)
-
-from websockets.sync.server import serve
-
 import torch
 import numpy as np
-
+from websockets.sync.server import serve
-from whisper_live.vad import VoiceActivityDetection
-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:
+except Exception:
     pass
 
+logging.basicConfig(level=logging.INFO)
+
+
+class VoiceActivityDetector:
+    def __init__(self, threshold=0.5):
+        self.model = VoiceActivityDetection()
+        self.threshold = threshold
+
+    def __call__(self, audio_frame):
+        speech_prob = self.model(torch.from_numpy(audio_frame), TranscriptionServer.RATE).item()
+        return speech_prob > self.threshold
+
+
+class ClientManager:
+    def __init__(self, max_clients=4, max_connection_time=600):
+        self.clients = {}
+        self.start_times = {}
+        self.max_clients = max_clients
+        self.max_connection_time = max_connection_time
+
+    def add_client(self, websocket, client):
+        self.clients[websocket] = client
+        self.start_times[websocket] = time.time()
+
+    def get_client(self, websocket):
+        if websocket in self.clients:
+            return self.clients[websocket]
+        return False
+
+    def remove_client(self, websocket):
+        client = self.clients.pop(websocket, None)
+        if client:
+            client.cleanup()
+        self.start_times.pop(websocket, None)
+
+    def get_wait_time(self):
+        """Calculate and return the estimated wait time for clients."""
+        wait_time = None
+        for start_time in self.start_times.values():
+            current_client_time_remaining = self.max_connection_time - (time.time() - start_time)
+            if wait_time is None or current_client_time_remaining < wait_time:
+                wait_time = current_client_time_remaining
+        return wait_time / 60 if wait_time is not None else 0
+
+    def is_server_full(self, websocket, options):
+        """Check if the server is full and send wait message if necessary."""
+        if len(self.clients) >= self.max_clients:
+            wait_time = self.get_wait_time()
+            response = {"uid": options["uid"], "status": "WAIT", "message": wait_time}
+            websocket.send(json.dumps(response))
+            return True
+        return False
+
+    def is_client_timeout(self, websocket):
+        elapsed_time = time.time() - self.start_times[websocket]
+        if elapsed_time >= self.max_connection_time:
+            self.clients[websocket].disconnect()
+            logging.warning(f"Client with uid '{self.clients[websocket].client_uid}' disconnected due to overtime.")
+            return True
+        return False
+
 
 class TranscriptionServer:
     """
@@ -42,12 +97,8 @@ class TranscriptionServer:
 
     def __init__(self):
         # voice activity detection model
-        
+        self.client_manager = ClientManager()
-        self.clients = {}
+        self.no_voice_activity_chunks = 0
-        self.websockets = {}
-        self.clients_start_time = {}
-        self.max_clients = 4
-        self.max_connection_time = 600
 
     def get_wait_time(self):
         """
@@ -58,7 +109,7 @@ class TranscriptionServer:
         """
         wait_time = None
 
-        for k, v in self.clients_start_time.items():
+        for _, v in self.clients_start_time.items():
             current_client_time_remaining = self.max_connection_time - (time.time() - v)
 
             if wait_time is None or current_client_time_remaining < wait_time:
@@ -66,6 +117,64 @@ class TranscriptionServer:
 
         return wait_time / 60
 
+    def is_server_full(self, websocket, options):
+        if len(self.clients) >= self.max_clients:
+            wait_time = self.get_wait_time()
+            response = {"uid": options["uid"], "status": "WAIT", "message": wait_time}
+            websocket.send(json.dumps(response))
+            websocket.close()
+            return True
+        return False
+
+    def initialize_client(
+        self, websocket, options, faster_whisper_custom_model_path,
+        whisper_tensorrt_path, trt_multilingual
+    ):
+        if self.backend == "tensorrt":
+            try:
+                client = ServeClientTensorRT(
+                    websocket,
+                    multilingual=trt_multilingual,
+                    language=options["language"],
+                    task=options["task"],
+                    client_uid=options["uid"],
+                    model=whisper_tensorrt_path
+                )
+                logging.info("Running TensorRT backend.")
+            except Exception as e:
+                logging.error(f"TensorRT-LLM not supported: {e}")
+                self.client_uid = options["uid"]
+                websocket.send(json.dumps({
+                    "uid": self.client_uid,
+                    "status": "WARNING",
+                    "message": "TensorRT-LLM not supported on Server yet. "
+                               "Reverting to available backend: 'faster_whisper'"
+                }))
+                self.backend = "faster_whisper"
+
+        if self.backend == "faster_whisper":
+            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("Running faster_whisper backend.")
+
+        # self.clients[websocket] = client
+        # self.clients_start_time[websocket] = time.time()
+        self.client_manager.add_client(websocket, client)
+
+    def get_audio_from_websocket(self, websocket):
+        frame_data = websocket.recv()
+        return np.frombuffer(frame_data, dtype=np.float32)
+
     def recv_audio(self,
                    websocket,
                    backend="faster_whisper",
@@ -96,127 +205,53 @@ class TranscriptionServer:
         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)
 
-        if len(self.clients) >= self.max_clients:
+        if self.client_manager.is_server_full(websocket, options):
-            logging.warning("Client Queue Full. Asking client to wait ...")
-            wait_time = self.get_wait_time()
-            response = {
-                "uid": options["uid"],
-                "status": "WAIT",
-                "message": wait_time,
-            }
-            websocket.send(json.dumps(response))
             websocket.close()
-            del websocket
             return
 
+        self.backend = backend
         if self.backend == "tensorrt":
+            self.vad_detector = VoiceActivityDetector()
+
+        self.initialize_client(
+            websocket, options, faster_whisper_custom_model_path, whisper_tensorrt_path, trt_multilingual)
+
+        while not self.client_manager.is_client_timeout(websocket):
             try:
-                import tensorrt as trt
+                frame_np = self.get_audio_from_websocket(websocket)
-                import tensorrt_llm
+                client = self.client_manager.get_client(websocket)
-                self.backend = "tensorrt"
-                client = ServeClientTensorRT(
-                    websocket,
-                    multilingual=trt_multilingual,
-                    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:
+                    if not self.voice_activity(websocket, frame_np):
-                        speech_prob = self.vad_model(torch.from_numpy(frame_np.copy()), self.RATE).item()
+                        continue
-                        if speech_prob < self.vad_threshold:
+                    self.no_voice_activity_chunks = 0
-                            no_voice_activity_chunks += 1
+                    client.set_eos(False)
-                            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:
+                client.add_frames(frame_np)
-                        logging.error(e)
-                        return
-                
-                self.clients[websocket].add_frames(frame_np)
-
-                elapsed_time = time.time() - self.clients_start_time[websocket]
-                if elapsed_time >= self.max_connection_time:
-                    self.clients[websocket].disconnect()
-                    logging.warning(f"Client with uid '{self.clients[websocket].client_uid}' disconnected due to overtime.")
-                    self.clients[websocket].cleanup()
-                    self.clients.pop(websocket)
-                    self.clients_start_time.pop(websocket)
-                    websocket.close()
-                    del websocket
-                    break
 
             except Exception as e:
                 logging.error(e)
-                self.clients[websocket].cleanup()
+                self.cleanup(websocket)
-                self.clients.pop(websocket)
+                websocket.close()
-                self.clients_start_time.pop(websocket)
-                del websocket
                 break
 
+        if self.client_manager.get_client(websocket):
+            self.cleanup(websocket)
+            websocket.close()
+        del websocket
+
     def run(self,
             host,
             port=9090,
             backend="tensorrt",
             faster_whisper_custom_model_path=None,
             whisper_tensorrt_path=None,
-            trt_multilingual=False
+            trt_multilingual=False):
-        ):
         """
         Run the transcription server.
 
@@ -237,6 +272,21 @@ class TranscriptionServer:
         ) as server:
             server.serve_forever()
 
+    def voice_activity(self, websocket, frame_np):
+        if not self.vad_detector(frame_np):
+            self.no_voice_activity_chunks += 1
+            if self.no_voice_activity_chunks > 3:
+                client = self.client_manager.get_client(websocket)
+                if not client.eos:
+                    client.set_eos(True)
+                time.sleep(0.1)    # Sleep 100m; wait some voice activity.
+            return False
+        return True
+
+    def cleanup(self, websocket):
+        if self.client_manager.get_client(websocket):
+            self.client_manager.remove_client(websocket)
+
 
 class ServeClientBase(object):
     RATE = 16000
@@ -254,7 +304,7 @@ class ServeClientBase(object):
         self.text = []
         self.current_out = ''
         self.prev_out = ''
-        self.t_start=None
+        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
@@ -269,6 +319,15 @@ class ServeClientBase(object):
         # threading
         self.lock = threading.Lock()
 
+    def speech_to_text(self):
+        raise NotImplementedError
+
+    def transcribe_audio(self):
+        raise NotImplementedError
+
+    def handle_transcription_output(self):
+        raise NotImplementedError
+
     def add_frames(self, frame_np):
         """
         Add audio frames to the ongoing audio stream buffer.
@@ -295,8 +354,49 @@ class ServeClientBase(object):
             self.frames_np = np.concatenate((self.frames_np, frame_np), axis=0)
         self.lock.release()
 
-    def speech_to_text(self):
+    def clip_audio_if_no_valid_segment(self):
-        raise NotImplementedError("Please implement in child Class.")
+        """
+        Update the timestamp offset based on audio buffer status.
+        Clip audio if the current chunk exceeds 30 seconds, this basically implies that
+        no valid segment for the last 30 seconds from whisper
+        """
+        if self.frames_np[int((self.timestamp_offset - self.frames_offset)*self.RATE):].shape[0] > 25 * self.RATE:
+            duration = self.frames_np.shape[0] / self.RATE
+            self.timestamp_offset = self.frames_offset + duration - 5
+
+    def get_audio_chunk_for_processing(self):
+        """Retrieve the next chunk of audio data for processing."""
+        samples_take = max(0, (self.timestamp_offset - self.frames_offset) * self.RATE)
+        input_bytes = self.frames_np[int(samples_take):].copy()
+        duration = input_bytes.shape[0] / self.RATE
+        return input_bytes, duration
+
+    def prepare_segments(self, last_segment=None):
+        """Prepare the segments to be sent to the client."""
+        segments = []
+        if len(self.transcript) >= self.send_last_n_segments:
+            segments = self.transcript[-self.send_last_n_segments:].copy()
+        else:
+            segments = self.transcript.copy()
+        if last_segment is not None:
+            segments = segments + [last_segment]
+        return segments
+
+    def get_audio_chunk_duration(self, input_bytes):
+        """Calculate the duration of the current audio chunk."""
+        return input_bytes.shape[0] / self.RATE
+
+    def send_transcription_to_client(self, segments):
+        """Send the transcription segments to the client."""
+        try:
+            self.websocket.send(
+                json.dumps({
+                    "uid": self.client_uid,
+                    "segments": segments,
+                })
+            )
+        except Exception as e:
+            logging.error(f"[ERROR]: Sending data to client: {e}")
 
     def disconnect(self):
         """
@@ -306,14 +406,10 @@ class ServeClientBase(object):
         that the transcription service is disconnecting gracefully.
 
         """
-        self.websocket.send(
+        self.websocket.send(json.dumps({
-            json.dumps(
+            "uid": self.client_uid,
-                {
+            "message": self.DISCONNECT
-                    "uid": self.client_uid,
+        }))
-                    "message": self.DISCONNECT
-                }
-            )
-        )
 
     def cleanup(self):
         """
@@ -357,16 +453,7 @@ class ServeClientTensorRT(ServeClientBase):
         pick_previous_segments (int): Number of previous segments to include in the output.
         websocket: The WebSocket connection for the client.
     """
-    def __init__(
+    def __init__(self, websocket, task="transcribe", multilingual=False, language=None, client_uid=None, model=None):
-        self,
-        websocket,
-        task="transcribe",
-        device=None,
-        multilingual=False,
-        language=None, 
-        client_uid=None,
-        model=None
-        ):
         """
         Initialize a ServeClient instance.
         The Whisper model is initialized based on the client's language and device availability.
@@ -400,15 +487,11 @@ class ServeClientTensorRT(ServeClientBase):
         self.trans_thread = threading.Thread(target=self.speech_to_text)
         self.trans_thread.start()
 
-        self.websocket.send(
+        self.websocket.send(json.dumps({
-            json.dumps(
+            "uid": self.client_uid,
-                {
+            "message": self.SERVER_READY,
-                    "uid": self.client_uid,
+            "backend": "tensorrt"
-                    "message": self.SERVER_READY,
+        }))
-                    "backend": "tensorrt"
-                }
-            )
-        )
 
     def warmup(self, warmup_steps=10):
         logging.info("[INFO:] Warming up TensorRT engine..")
@@ -421,31 +504,27 @@ class ServeClientTensorRT(ServeClientBase):
         self.eos = eos
         self.lock.release()
 
-    def add_frames(self, frame_np):
+    def handle_transcription_output(self, last_segment, duration):
-        """
+        """Handle the transcription output, updating the transcript and sending data to the client."""
-        Add audio frames to the ongoing audio stream buffer.
+        segments = self.prepare_segments({"text": last_segment})
+        self.send_transcription_to_client(segments)
+        if self.eos:
+            self.update_timestamp_offset(last_segment, duration)
 
-        This method is responsible for maintaining the audio stream buffer, allowing the continuous addition
+    def transcribe_audio(self, input_bytes):
-        of audio frames as they are received. It also ensures that the buffer does not exceed a specified size
+        """Transcribe the audio chunk and send the results to the client."""
-        to prevent excessive memory usage.
+        logging.info(f"[WhisperTensorRT:] Processing audio with duration: {input_bytes.shape[0] / self.RATE}")
+        mel, duration = self.transcriber.log_mel_spectrogram(input_bytes)
+        last_segment = self.transcriber.transcribe(mel)
+        if last_segment:
+            self.handle_transcription_output(last_segment, duration)
 
-        If the buffer size exceeds a threshold (45 seconds of audio data), it discards the oldest 30 seconds
+    def update_timestamp_offset(self, last_segment, duration):
-        of audio data to maintain a reasonable buffer size. If the buffer is empty, it initializes it with the provided
+        if not len(self.transcript):
-        audio frame. The audio stream buffer is used for real-time processing of audio data for transcription.
+            self.transcript.append({"text": last_segment + " "})
-
+        elif self.transcript[-1]["text"].strip() != last_segment:
-        Args:
+            self.transcript.append({"text": last_segment + " "})
-            frame_np (numpy.ndarray): The audio frame data as a NumPy array.
+        self.timestamp_offset += duration
-
-        """
-        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):
         """
@@ -473,49 +552,16 @@ class ServeClientTensorRT(ServeClientBase):
                 time.sleep(0.02)    # wait for any audio to arrive
                 continue
 
-            # clip audio if the current chunk exceeds 30 seconds, this basically implies that
+            self.clip_audio_if_no_valid_segment()
-            # 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, duration = self.get_audio_chunk_for_processing()
-            input_bytes = self.frames_np[int(samples_take):].copy()
+            if duration < 0.4:
-            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)
+                self.transcribe_audio(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}")
@@ -550,17 +596,8 @@ class ServeClientFasterWhisper(ServeClientBase):
         pick_previous_segments (int): Number of previous segments to include in the output.
         websocket: The WebSocket connection for the client.
     """
-    def __init__(
+    def __init__(self, websocket, task="transcribe", device=None, language=None, client_uid=None, model="small.en",
-        self,
+                 initial_prompt=None, vad_parameters=None):
-        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.
@@ -592,13 +629,13 @@ class ServeClientFasterWhisper(ServeClientBase):
 
         device = "cuda" if torch.cuda.is_available() else "cpu"
 
-        if self.model_size_or_path == None:
+        if self.model_size_or_path is None:
             return
 
         self.transcriber = WhisperModel(
             self.model_size_or_path,
             device=device,
-            compute_type="int8" if device=="cpu" else "float16", 
+            compute_type="int8" if device == "cpu" else "float16",
             local_files_only=False,
         )
 
@@ -638,6 +675,38 @@ class ServeClientFasterWhisper(ServeClientBase):
             return None
         return model_size
 
+    def set_language(self, info):
+        if info.language_probability > 0.5:
+            self.language = info.language
+            logging.info(f"Detected language {self.language} with probability {info.language_probability}")
+            self.websocket.send(json.dumps(
+                {"uid": self.client_uid, "language": self.language, "language_prob": info.language_probability}))
+
+    def transcribe_audio(self, input_sample):
+        result, info = self.transcriber.transcribe(
+            input_sample,
+            initial_prompt=self.initial_prompt,
+            language=self.language,
+            task=self.task,
+            vad_filter=True,
+            vad_parameters=self.vad_parameters)
+        if self.language is None:
+            self.set_language(info)
+        return result
+
+    def get_previous_output(self):
+        segments = []
+        if self.t_start is None:
+            self.t_start = time.time()
+        if time.time() - self.t_start < self.show_prev_out_thresh:
+            segments = self.prepare_segments()
+
+        # add a blank if there is no speech for 3 seconds
+        if len(self.text) and self.text[-1] != '':
+            if time.time() - self.t_start > self.add_pause_thresh:
+                self.text.append('')
+        return segments
+
     def speech_to_text(self):
         """
         Process an audio stream in an infinite loop, continuously transcribing the speech.
@@ -663,74 +732,29 @@ class ServeClientFasterWhisper(ServeClientBase):
             if self.frames_np is None:
                 continue
 
-            # clip audio if the current chunk exceeds 30 seconds, this basically implies that
+            self.clip_audio_if_no_valid_segment()
-            # 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, duration = self.get_audio_chunk_for_processing()
-            input_bytes = self.frames_np[int(samples_take):].copy()
+            if duration < 1.0:
-            duration = input_bytes.shape[0] / self.RATE
-            if duration<1.0:
                 continue
             try:
                 input_sample = input_bytes.copy()
-                
+                result = self.transcribe_audio(input_sample)
-                # whisper transcribe with prompt
-                result, info = self.transcriber.transcribe(
-                    input_sample, 
-                    initial_prompt=self.initial_prompt,
-                    language=self.language,
-                    task=self.task,
-                    vad_filter=True,
-                    vad_parameters=self.vad_parameters
-                )
 
                 if self.language is None:
-                    if info.language_probability > 0.5:
+                    continue
-                        self.language = info.language
-                        logging.info(f"Detected language {self.language} with probability {info.language_probability}")
-                        self.websocket.send(json.dumps(
-                            {"uid": self.client_uid, "language": self.language, "language_prob": info.language_probability}))
-                    else:
-                        # detect language again
-                        continue
 
                 if len(result):
                     self.t_start = None
                     last_segment = self.update_segments(result, duration)
-                    if len(self.transcript) < self.send_last_n_segments:
+                    segments = self.prepare_segments(last_segment)
-                        segments = self.transcript
-                    else:
-                        segments = self.transcript[-self.send_last_n_segments:]
-                    if last_segment is not None:
-                        segments = segments + [last_segment]                    
                 else:
                     # show previous output if there is pause i.e. no output from whisper
-                    segments = []
+                    segments = self.get_previous_output()
-                    if self.t_start is None: self.t_start = time.time()
-                    if time.time() - self.t_start < self.show_prev_out_thresh:
-                        if len(self.transcript) < self.send_last_n_segments:
-                            segments = self.transcript
-                        else:
-                            segments = self.transcript[-self.send_last_n_segments:]
 
-                    # add a blank if there is no speech for 3 seconds
+                if not len(segments):
-                    if len(self.text) and self.text[-1] != '':
+                    continue
-                        if time.time() - self.t_start > self.add_pause_thresh:
+                self.send_transcription_to_client(segments)
-                            self.text.append('')
-
-                if not len(segments): continue
-                try:
-                    self.websocket.send(
-                        json.dumps({
-                            "uid": self.client_uid,
-                            "segments": segments
-                        })
-                    )
-                except Exception as e:
-                    logging.error(f"[ERROR]: Failed to send message to client: {e}")
 
             except Exception as e:
                 logging.error(f"[ERROR]: Failed to transcribe audio chunk: {e}")
@@ -775,11 +799,12 @@ class ServeClientFasterWhisper(ServeClientBase):
                 self.text.append(text_)
                 start, end = self.timestamp_offset + s.start, self.timestamp_offset + min(duration, s.end)
 
-                if start >= end: continue
+                if start >= end:
-                if s.no_speech_prob > self.no_speech_thresh: continue
+                    continue
+                if s.no_speech_prob > self.no_speech_thresh:
+                    continue
 
                 self.transcript.append(self.format_segment(start, end, text_))
-                
                 offset = min(duration, s.end)
 
         self.current_out += segments[-1].text
@@ -797,7 +822,7 @@ class ServeClientFasterWhisper(ServeClientBase):
             self.same_output_threshold = 0
 
         if self.same_output_threshold > 5:
-            if not len(self.text) or self.text[-1].strip().lower()!=self.current_out.strip().lower():          
+            if not len(self.text) or self.text[-1].strip().lower() != self.current_out.strip().lower():
                 self.text.append(self.current_out)
                 self.transcript.append(self.format_segment(
                     self.timestamp_offset,

whisper_live/tensorrt_utils.py

@@ -214,7 +214,7 @@ def store_transcripts(filename: Pathlike, texts: Iterable[Tuple[str, str,
             print(f"{cut_id}:\thyp={hyp}", file=f)
 
 
-def write_error_stats(
+def write_error_stats(                                              # noqa: C901
     f: TextIO,
     test_set_name: str,
     results: List[Tuple[str, str]],

whisper_live/transcriber.py

@@ -400,7 +400,7 @@ class WhisperModel:
 
         return segments, info
 
-    def generate_segments(
+    def generate_segments(                                                  # noqa: C901
         self,
         features: np.ndarray,
         tokenizer: Tokenizer,
@@ -425,7 +425,7 @@ class WhisperModel:
         all_segments = []
         while seek < content_frames:
             time_offset = seek * self.feature_extractor.time_per_frame
-            segment = features[:, seek : seek + self.feature_extractor.nb_max_frames]
+            segment = features[:, seek:seek + self.feature_extractor.nb_max_frames]
             segment_size = min(
                 self.feature_extractor.nb_max_frames, content_frames - seek
             )
@@ -749,7 +749,7 @@ class WhisperModel:
 
         if previous_tokens:
             prompt.append(tokenizer.sot_prev)
-            prompt.extend(previous_tokens[-(self.max_length // 2 - 1) :])
+            prompt.extend(previous_tokens[-(self.max_length // 2 - 1):])
 
         prompt.extend(tokenizer.sot_sequence)
 
@@ -766,7 +766,7 @@ class WhisperModel:
 
         return prompt
 
-    def add_word_timestamps(
+    def add_word_timestamps(                                                    # noqa: C901
         self,
         segments: List[dict],
         tokenizer: Tokenizer,

whisper_live/transcriber_tensorrt.py

@@ -1,17 +1,14 @@
-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
+from typing import Union
 
 import torch
 import numpy as np
+import torch.nn.functional as F
 from whisper.tokenizer import get_tokenizer
-from whisper_live.tensorrt_utils import (mel_filters, store_transcripts,
+from whisper_live.tensorrt_utils import (mel_filters, load_audio_wav_format, pad_or_trim, load_audio)
-                           write_error_stats, load_audio_wav_format,
-                           pad_or_trim, load_audio)
 
 import tensorrt_llm
 import tensorrt_llm.logger as logger
@@ -38,8 +35,6 @@ class WhisperEncoding:
         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']
@@ -176,16 +171,8 @@ class WhisperDecoding:
 
 class WhisperTRTLLM(object):
 
-    def __init__(
+    def __init__(self, engine_dir, assets_dir=None, device=None, is_multilingual=False,
-        self,
+                 language="en", task="transcribe"):
-        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)
@@ -212,7 +199,7 @@ class WhisperTRTLLM(object):
         self,
         audio: Union[str, np.ndarray, torch.Tensor],
         padding: int = 0,
-        return_duration = True
+        return_duration=True
     ):
         """
         Compute the log-Mel spectrogram of
@@ -242,8 +229,7 @@ class WhisperTRTLLM(object):
                     audio, _ = load_audio_wav_format(audio)
                 else:
                     audio = load_audio(audio)
-            assert isinstance(audio,
+            assert isinstance(audio, np.ndarray), f"Unsupported audio type: {type(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)
@@ -254,14 +240,9 @@ class WhisperTRTLLM(object):
         if padding > 0:
             audio = F.pad(audio, (0, padding))
         window = torch.hann_window(N_FFT).to(audio.device)
-        stft = torch.stft(audio,
+        stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)
-                        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()
@@ -272,7 +253,6 @@ class WhisperTRTLLM(object):
         else:
             return log_spec
 
-
     def process_batch(
             self,
             mel,

whisper_live/utils.py

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

whisper_live/vad.py

@@ -34,7 +34,7 @@ class VoiceActivityDetection():
 
         if sr != 16000 and (sr % 16000 == 0):
             step = sr // 16000
-            x = x[:,::step]
+            x = x[:, ::step]
             sr = 16000
 
         if sr not in self.sample_rates: