From 735d6c77631c246f5a766aeec4b4a4d2cb42b124 Mon Sep 17 00:00:00 2001 From: makaveli10 Date: Fri, 19 Jan 2024 11:43:21 +0000 Subject: [PATCH] merge with main --- run_server.py | 28 +++++++++++------ whisper_live/server.py | 71 +++++++++++++++++++++++++++--------------- 2 files changed, 65 insertions(+), 34 deletions(-) diff --git a/run_server.py b/run_server.py index 6ba1c66..3feeec6 100644 --- a/run_server.py +++ b/run_server.py @@ -3,26 +3,36 @@ from whisper_live.server import TranscriptionServer if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--port', type=int, default=9090, help="Websocket port to run the server on.") - parser.add_argument('--backend', type=str, default='faster_whisper', help='Backends from ["tensorrt", "faster_whisper"]') - parser.add_argument('--whisper_tensorrt_path', + 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', + 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.whisper_tensorrt_path is None: + if args.trt_model_path is None: raise ValueError("Please Provide a valid tensorrt model path") server = TranscriptionServer() server.run( "0.0.0.0", - port=6006, - backend=args.backend, - whisper_tensorrt_path=args.whisper_tensorrt_path, - multilingual=args.trt_multilingual + 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 ) diff --git a/whisper_live/server.py b/whisper_live/server.py index 5338050..cf6f8d0 100644 --- a/whisper_live/server.py +++ b/whisper_live/server.py @@ -1,3 +1,4 @@ +import os import websockets import time import threading @@ -68,7 +69,12 @@ class TranscriptionServer: return wait_time / 60 - def recv_audio(self, websocket, backend="tensorrt", whisper_tensorrt_path=None, multilingual=False): + 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. @@ -85,7 +91,11 @@ 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. """ @@ -118,11 +128,11 @@ class TranscriptionServer: self.backend = "tensorrt" client = ServeClientTensorRT( websocket, - multilingual=multilingual, + multilingual=trt_multilingual, language=options["language"], task=options["task"], client_uid=options["uid"], - model_path=whisper_tensorrt_path + model=whisper_tensorrt_path ) logging.info(f"Running TensorRT backend.") except Exception as e: @@ -138,13 +148,17 @@ class TranscriptionServer: 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, multilingual=options["multilingual"], language=options["language"], task=options["task"], client_uid=options["uid"], - model_size=options["model_size"], + model=options["model"], initial_prompt=options.get("initial_prompt"), vad_parameters=options.get("vad_parameters") ) @@ -159,7 +173,7 @@ class TranscriptionServer: frame_data = websocket.recv() frame_np = np.frombuffer(frame_data, dtype=np.float32) - # VAD + # 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() @@ -168,7 +182,7 @@ class TranscriptionServer: if no_voice_activity_chunks > 3: if not self.clients[websocket].eos: self.clients[websocket].set_eos(True) - time.sleep(0.1) # EOS stop receiving frames for a 100ms(to send output to LLM.) + time.sleep(0.1) # Sleep 100m; wait some voice activity. continue no_voice_activity_chunks = 0 self.clients[websocket].set_eos(False) @@ -198,7 +212,14 @@ class TranscriptionServer: del websocket break - def run(self, host, port=9090, backend="tensorrt", whisper_tensorrt_path=None, multilingual=False): + 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. @@ -210,8 +231,9 @@ class TranscriptionServer: functools.partial( self.recv_audio, backend=backend, + faster_whisper_custom_model_path=faster_whisper_custom_model_path, whisper_tensorrt_path=whisper_tensorrt_path, - multilingual=multilingual + trt_multilingual=trt_multilingual ), host, port @@ -220,6 +242,10 @@ class TranscriptionServer: 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 @@ -303,7 +329,6 @@ class ServeClientBase(object): """ logging.info("Cleaning up.") self.exit = True - self.transcriber.destroy() class ServeClientTensorRT(ServeClientBase): @@ -335,10 +360,6 @@ class ServeClientTensorRT(ServeClientBase): 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, @@ -347,7 +368,7 @@ class ServeClientTensorRT(ServeClientBase): multilingual=False, language=None, client_uid=None, - model_path=None + model=None ): """ Initialize a ServeClient instance. @@ -369,7 +390,7 @@ class ServeClientTensorRT(ServeClientBase): self.task = task self.eos = False self.transcriber = WhisperTRTLLM( - model_path, + model, assets_dir="assets", device="cuda", is_multilingual=multilingual, @@ -530,10 +551,6 @@ class ServeClientFasterWhisper(ServeClientBase): 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, @@ -542,7 +559,7 @@ class ServeClientFasterWhisper(ServeClientBase): multilingual=False, language=None, client_uid=None, - model_size="small", + model="small", initial_prompt=None, vad_parameters=None ): @@ -563,10 +580,14 @@ class ServeClientFasterWhisper(ServeClientBase): """ super().__init__(client_uid, websocket) self.model_sizes = [ - "tiny", "base", "small", "medium", "large-v2", "large-v3" + "tiny", "tiny.en", "base", "base.en", "small", "small.en", + "medium", "medium.en", "large-v2", "large-v3", ] self.multilingual = multilingual - self.model_size = self.get_model_size(model_size) + if not os.path.exists(model): + self.model_size_or_path = self.get_model_size(model) + else: + self.model_size_or_path = model self.language = language if self.multilingual else "en" self.task = task self.initial_prompt = initial_prompt @@ -574,11 +595,11 @@ class ServeClientFasterWhisper(ServeClientBase): device = "cuda" if torch.cuda.is_available() else "cpu" - if self.model_size == None: + if self.model_size_or_path == None: return self.transcriber = WhisperModel( - self.model_size, + self.model_size_or_path, device=device, compute_type="int8" if device=="cpu" else "float16", local_files_only=False,