add whisper_live module

whisper_live/client.py

@@ -12,204 +12,6 @@ import json
 import websocket
 
 
-CHUNK = 1024
-FORMAT = pyaudio.paInt16
-CHANNELS = 1
-RATE = 16000
-RECORD_SECONDS = 60000
-START_RECORDING = False
-multilingual = False
-language = None
-
-
-
-def on_message(ws, message):
-    global START_RECORDING, language
-    message = json.loads(message)
-    if message == "SERVER_READY":
-        START_RECORDING = True
-        return
-
-    if isinstance(message, dict):
-        language = message.get("language")
-        lang_prob = message.get("language_prob")
-        print(f"Server detected language {language} with probability {lang_prob}")
-        return
-
-    text = []
-    if len(message):
-        for seg in message:
-            if len(text):
-                if text[-1] != seg["text"]:
-                    text.append(seg["text"])
-            else:
-                text.append(seg["text"])
-    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(ws, error):
-    print(error)
-
-def on_close(ws, close_status_code, close_msg):
-    print("### websocket connection closed ###")
-
-def on_open(ws):
-    global multilingual, language, task
-    print(multilingual, language, task)
-
-    print("Opened connection")
-    ws.send(json.dumps({
-        'multilingual': multilingual[0],
-        'language': language[0],
-        'task': task
-    }))
-
-
-class Client:
-    def __init__(self, host=None, port=None):
-        self.timestamp_offset = 0.0
-        self.audio_bytes = None
-        self.p = pyaudio.PyAudio()
-        self.stream = self.p.open(format=FORMAT,
-                        channels=CHANNELS,
-                        rate=RATE,
-                        input=True,
-                        frames_per_buffer=CHUNK)
-        print(self.p.get_sample_size(FORMAT))
-
-        # create websocket connection
-        if host is not None and port is not None:
-            socket_url = f"ws://{host}:{port}"    
-            self.client_socket = websocket.WebSocketApp(socket_url,
-                                  on_open=on_open,
-                                  on_message=on_message,
-                                  on_error=on_error,
-                                  on_close=on_close)
-        else:
-            print("No host or port specified.")
-            return
-
-        # start websocket client in a thread
-        self.ws_thread = threading.Thread(target=self.client_socket.run_forever)
-        self.ws_thread.setDaemon(True)
-        self.ws_thread.start()
-
-        self.frames = b""
-        print("* recording")
-
-    def send_packet_to_server(self, message):
-        try:
-            self.client_socket.send(message, websocket.ABNF.OPCODE_BINARY)
-        except Exception as e:
-            print(e)
-
-    @staticmethod
-    def bytes_to_float_array(audio_bytes):
-        raw_data = np.frombuffer(
-            buffer=audio_bytes, dtype=np.int16
-        )
-        return raw_data.astype(np.float32) / 32768.0
-    
-    def play_file(self, filename):
-        # read audio and create pyaudio stream
-        self.wf = wave.open(filename, 'rb')
-        self.stream = self.p.open(format=self.p.get_format_from_width(self.wf.getsampwidth()),
-                channels=self.wf.getnchannels(),
-                rate=self.wf.getframerate(),
-                input=True,
-                output=True,
-                frames_per_buffer=CHUNK)
-        try:
-            while True:
-                data = self.wf.readframes(CHUNK)
-                if data==b'': break
-
-                audio_array = Client.bytes_to_float_array(data)
-                self.send_packet_to_server(audio_array.tobytes())
-                self.stream.write(data)
-
-            self.wf.close()
-            self.stream.close()
-
-        except KeyboardInterrupt:
-            print("Keyboard interrupt.")
-
-
-    def get_client_socket(self):
-        return self.client_socket
-    
-    def write_audio_frames_to_file(self, frames, file_name):
-        wf = wave.open(file_name, 'wb')
-        wf.setnchannels(CHANNELS)
-        wf.setsampwidth(2)
-        wf.setframerate(RATE)
-        wf.writeframes(frames)
-        wf.close()
-
-    def record(self, out_file="output_recording.wav"):
-        n_audio_file = 0
-        # create dir for saving audio chunks
-        if not os.path.exists("chunks"):
-            os.makedirs("chunks", exist_ok=True)
-        try:
-            for _ in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
-                data = self.stream.read(CHUNK)
-                self.frames += data
-
-                audio_array = Client.bytes_to_float_array(data)
-                
-                self.send_packet_to_server(audio_array.tobytes())
-
-                # save frames if more than a minute
-                if len(self.frames) > 60*RATE:
-                    t = threading.Thread(
-                        target=self.write_audio_frames_to_file,
-                        args=(self.frames[:], f"chunks/{n_audio_file}.wav", )
-                    )
-                    t.start()
-                    n_audio_file += 1
-                    self.frames = b""
-
-        except KeyboardInterrupt:
-            if len(self.frames):
-                self.write_audio_frames_to_file(
-                    self.frames[:], f"chunks/{n_audio_file}.wav")
-                n_audio_file += 1
-            self.stream.stop_stream()
-            self.stream.close()
-            self.p.terminate()
-
-
-            # combine all the audio files
-            self.write_output_recording(n_audio_file, out_file)
-    
-    def write_output_recording(self, n_audio_file, out_file):
-        input_files = [f"chunks/{i}.wav" for i in range(n_audio_file) if os.path.exists(f"chunks/{i}.wav")]
-        wf = wave.open(out_file, 'wb')
-        wf.setnchannels(CHANNELS)
-        wf.setsampwidth(2)
-        wf.setframerate(RATE)
-        for in_file in input_files:
-            w = wave.open(in_file, 'rb')
-            while True:
-                data = w.readframes(CHUNK)
-                if data==b'': break
-                wf.writeframes(data)
-            w.close()
-            # remove this file
-            os.remove(in_file)
-        wf.close()
-
-
 def resample(file: str, sr: int = 16000):
     """
     # https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/audio.py#L22
@@ -239,30 +41,235 @@ def resample(file: str, sr: int = 16000):
     return resampled_file
 
 
-if __name__=="__main__":
+class Client:
-    parser = argparse.ArgumentParser()
+    CHUNK = 1024
-    parser.add_argument('--audio', type=str, help='audio file to transcribe')
+    FORMAT = pyaudio.paInt16
-    parser.add_argument('--host', default=None, type=str, help='websocket server address to connect to')
+    CHANNELS = 1
-    parser.add_argument('--port', default=None, type=str, help='websocket server port to connect to')
+    RATE = 16000
-    parser.add_argument('--multilingual', action="store_true", help='use multilingual model')
+    RECORD_SECONDS = 60000
-    parser.add_argument('--language', default=None, type=str, help='languages to use')
+    START_RECORDING = False
-    parser.add_argument(
+    multilingual = False
-        '--task', default="transcribe", type=str, help='task transcribe/translate (translates from any to english)')
+    language = None
-    opt = parser.parse_args()
+    task = "transcribe"
-    print(opt)
+    
-    multilingual=opt.multilingual,
+    def __init__(self, host=None, port=None, is_multilingual=False, lang=None, translate=False):
-    language = opt.language if opt.multilingual else "en",
+        Client.multilingual = is_multilingual
-    task = opt.task
+        Client.language = lang if is_multilingual else "en"
-    c = Client(host=opt.host, port=opt.port)
+        if translate:
+            Client.task = "translate"
 
-    # while loop to wait for server to be ready
+        self.timestamp_offset = 0.0
-    print("Waiting for server ready ...")
+        self.audio_bytes = None
-    while not START_RECORDING:
+        self.p = pyaudio.PyAudio()
-        pass
+        self.stream = self.p.open(format=self.FORMAT,
-    print("Server Ready!")
+                        channels=self.CHANNELS,
+                        rate=self.RATE,
+                        input=True,
+                        frames_per_buffer=self.CHUNK)
+
+        # create websocket connection
+        if host is not None and port is not None:
+            socket_url = f"ws://{host}:{port}"    
+            self.client_socket = websocket.WebSocketApp(socket_url,
+                                  on_open=Client.on_open,
+                                  on_message=Client.on_message,
+                                  on_error=Client.on_error,
+                                  on_close=Client.on_close)
+        else:
+            print("[ERROR]: No host or port specified.")
+            return
+
+        # start websocket client in a thread
+        self.ws_thread = threading.Thread(target=self.client_socket.run_forever)
+        self.ws_thread.setDaemon(True)
+        self.ws_thread.start()
+
+        self.frames = b""
+        print("[INFO]: * recording")
+    
+    @staticmethod
+    def on_message(ws, message):
+        message = json.loads(message)
+        if message == "SERVER_READY":
+            Client.START_RECORDING = True
+            return
+
+        if isinstance(message, dict):
+            Client.language = message.get("language")
+            lang_prob = message.get("language_prob")
+            print(f"[INFO]: Server detected language {Client.language} with probability {lang_prob}")
+            return
+
+        text = []
+        if len(message):
+            for seg in message:
+                if len(text):
+                    if text[-1] != seg["text"]:
+                        text.append(seg["text"])
+                else:
+                    text.append(seg["text"])
+        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)
+
+    @staticmethod
+    def on_error(ws, error):
+        print(error)
+
+    @staticmethod
+    def on_close(ws, close_status_code, close_msg):
+        print(f"[INFO]: Websocket connection closed.")
+
+    @staticmethod
+    def on_open(ws):
+        print(Client.multilingual, Client.language, Client.task)
+
+        print("[INFO]: Opened connection")
+        ws.send(json.dumps({
+            'multilingual': Client.multilingual,
+            'language': Client.language,
+            'task': Client.task
+        }))
+
+    @staticmethod
+    def bytes_to_float_array(audio_bytes):
+        raw_data = np.frombuffer(
+            buffer=audio_bytes, dtype=np.int16
+        )
+        return raw_data.astype(np.float32) / 32768.0
+    
+    def send_packet_to_server(self, message):
+        try:
+            self.client_socket.send(message, websocket.ABNF.OPCODE_BINARY)
+        except Exception as e:
+            print(e)
+    
+    def play_file(self, filename):
+        # read audio and create pyaudio stream
+        self.wf = wave.open(filename, 'rb')
+        self.stream = self.p.open(format=self.p.get_format_from_width(self.wf.getsampwidth()),
+                channels=self.wf.getnchannels(),
+                rate=self.wf.getframerate(),
+                input=True,
+                output=True,
+                frames_per_buffer=self.CHUNK)
+        try:
+            while True:
+                data = self.wf.readframes(self.CHUNK)
+                if data==b'': break
+
+                audio_array = Client.bytes_to_float_array(data)
+                self.send_packet_to_server(audio_array.tobytes())
+                self.stream.write(data)
+
+            self.wf.close()
+            self.stream.close()
+
+        except KeyboardInterrupt:
+            self.wf.close()
+            self.stream.stop_stream()
+            self.stream.close()
+            self.p.terminate()
+            self.close_websocket()
+            print("[INFO]: Keyboard interrupt.")
+
+    def close_websocket(self):
+        try:
+            self.client_socket.close()  # Close the WebSocket connection
+        except Exception as e:
+            print("[ERROR]: Error closing WebSocket:", e)
+
+        try:
+            self.ws_thread.join()  # Wait for the WebSocket thread to finish
+        except Exception as e:
+            print("[ERROR:] Error joining WebSocket thread:", e)
+
+    def get_client_socket(self):
+        return self.client_socket
+    
+    def write_audio_frames_to_file(self, frames, file_name):
+        wf = wave.open(file_name, 'wb')
+        wf.setnchannels(self.CHANNELS)
+        wf.setsampwidth(2)
+        wf.setframerate(self.RATE)
+        wf.writeframes(frames)
+        wf.close()
+
+    def record(self, out_file="output_recording.wav"):
+        n_audio_file = 0
+        # create dir for saving audio chunks
+        if not os.path.exists("chunks"):
+            os.makedirs("chunks", exist_ok=True)
+        try:
+            for _ in range(0, int(self.RATE / self.CHUNK * self.RECORD_SECONDS)):
+                data = self.stream.read(self.CHUNK)
+                self.frames += data
+
+                audio_array = Client.bytes_to_float_array(data)
+                
+                self.send_packet_to_server(audio_array.tobytes())
+
+                # save frames if more than a minute
+                if len(self.frames) > 60*self.RATE:
+                    t = threading.Thread(
+                        target=self.write_audio_frames_to_file,
+                        args=(self.frames[:], f"chunks/{n_audio_file}.wav", )
+                    )
+                    t.start()
+                    n_audio_file += 1
+                    self.frames = b""
+
+        except KeyboardInterrupt:
+            if len(self.frames):
+                self.write_audio_frames_to_file(
+                    self.frames[:], f"chunks/{n_audio_file}.wav")
+                n_audio_file += 1
+            self.stream.stop_stream()
+            self.stream.close()
+            self.p.terminate()
+            self.close_websocket()
+
+            # combine all the audio files
+            self.write_output_recording(n_audio_file, out_file)
+    
+    def write_output_recording(self, n_audio_file, out_file):
+        input_files = [f"chunks/{i}.wav" for i in range(n_audio_file) if os.path.exists(f"chunks/{i}.wav")]
+        wf = wave.open(out_file, 'wb')
+        wf.setnchannels(self.CHANNELS)
+        wf.setsampwidth(2)
+        wf.setframerate(self.RATE)
+        for in_file in input_files:
+            w = wave.open(in_file, 'rb')
+            while True:
+                data = w.readframes(self.CHUNK)
+                if data==b'': break
+                wf.writeframes(data)
+            w.close()
+            # remove this file
+            os.remove(in_file)
+        wf.close()
+
+
+class TranscriptionClient:
+    def __init__(self, host, port, is_multilingual=False, lang=None, translate=False):
+        self.client = Client(host, port, is_multilingual, lang, translate)
+        
+    def __call__(self, audio=None):
+        print("[INFO]: Waiting for server ready ...")
+        while not Client.START_RECORDING:
+            pass
+        print("[INFO]: Server Ready!")
+        if audio is not None:
+            resampled_file = resample(audio)
+            self.client.play_file(resampled_file)
+        else:
+            self.client.record()
 
-    if opt.audio is not None:
-        resampled_file = resample(opt.audio)
-        c.play_file(resampled_file)
-    else:
-        c.record()

whisper_live/server.py

@@ -10,48 +10,70 @@ logging.basicConfig(level = logging.INFO)
 
 from collections import deque
 from dataclasses import dataclass
+from websockets.sync.server import serve
 
 import torch
 import numpy as np
-from websockets.sync import server
+from whisper_live.transcriber import WhisperModel
-from websockets.sync.server import serve
-from transcriber import WhisperModel
 
 
-clients = {}
+class TranscriptionServer:
-SERVER_READY = "SERVER_READY"
-
-def recv_audio(websocket):
     """
-    Receive audio chunks from client in an infinite loop.
+    Represents a transcription server that handles incoming audio from clients.
-    """
-    global clients
-    options = websocket.recv()
-    options = json.loads(options)
-    client = ServeClient(
-        websocket, 
-        multilingual=options["multilingual"],
-        language=options["language"],
-        task=options["task"]
-    )
 
-    clients[websocket] = client
+    Attributes:
-    
+        clients (dict): A dictionary to store connected clients.
-    while True:
+    """
-        try:
+
-            frame_data = websocket.recv()
+    def __init__(self):
-            frame_np = np.frombuffer(frame_data, np.float32)
+        self.clients = {}
-            clients[websocket].add_frames(frame_np)
+
-            
+    def recv_audio(self, websocket):
-        except Exception as e:
+        """
-            clients[websocket].cleanup()
+        Receive audio chunks from a client in an infinite loop.
-            clients.pop(websocket)
+
-            logging.info("Connection Closed.")
+        Args:
-            break
+            websocket (WebSocket): The WebSocket connection for the client.
+        """
+        options = websocket.recv()
+        options = json.loads(options)
+        client = ServeClient(
+            websocket,
+            multilingual=options["multilingual"],
+            language=options["language"],
+            task=options["task"],
+        )
+
+        self.clients[websocket] = client
+
+        while True:
+            try:
+                frame_data = websocket.recv()
+                frame_np = np.frombuffer(frame_data, np.float32)
+                self.clients[websocket].add_frames(frame_np)
+
+            except Exception as e:
+                self.clients[websocket].cleanup()
+                self.clients.pop(websocket)
+                logging.info("Connection Closed.")
+                break
+
+    def run(self, host, port):
+        """
+        Run the transcription server.
+
+        Args:
+            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:
+            server.serve_forever()
 
 
 class ServeClient:
     RATE = 16000
+    SERVER_READY = "SERVER_READY"
+
     def __init__(self, websocket, task="transcribe", device=None, multilingual=False, language=None):
         self.data = b""
         self.frames = b""
@@ -94,7 +116,7 @@ class ServeClient:
         self.websocket = websocket
         self.trans_thread = threading.Thread(target=self.speech_to_text)
         self.trans_thread.start()
-        self.websocket.send(json.dumps(SERVER_READY))
+        self.websocket.send(json.dumps(self.SERVER_READY))
     
     def fill_output(self, output):
         """
@@ -302,8 +324,4 @@ class ServeClient:
         logging.info("Cleaning up.")
         self.exit = True
         self.transcriber.destroy()
-    
 
-if __name__ == "__main__":
-    with serve(recv_audio, "0.0.0.0", 9090) as server:
-        server.serve_forever()