diff --git a/whisper_live/server.py b/whisper_live/server.py index 18bf595..c06d0cd 100644 --- a/whisper_live/server.py +++ b/whisper_live/server.py @@ -101,9 +101,10 @@ class TranscriptionServer: multilingual=options["multilingual"], language=options["language"], task=options["task"], - client_uid=options["uid"] + client_uid=options["uid"], + model_size=options["model_size"] ) - + self.clients[websocket] = client self.clients_start_time[websocket] = time.time() @@ -127,7 +128,8 @@ class TranscriptionServer: except Exception as e: logging.error(e) - self.clients[websocket].cleanup() + if self.clients[websocket].model_size is not None: + self.clients[websocket].cleanup() self.clients.pop(websocket) self.clients_start_time.pop(websocket) logging.info("Connection Closed.") @@ -180,7 +182,16 @@ class ServeClient: SERVER_READY = "SERVER_READY" DISCONNECT = "DISCONNECT" - def __init__(self, websocket, task="transcribe", device=None, multilingual=False, language=None, client_uid=None): + def __init__( + self, + websocket, + task="transcribe", + device=None, + multilingual=False, + language=None, + client_uid=None, + model_size="small" + ): """ Initialize a ServeClient instance. The Whisper model is initialized based on the client's language and device availability. @@ -199,11 +210,23 @@ class ServeClient: self.client_uid = client_uid self.data = b"" self.frames = b"" - self.language = language if multilingual else "en" + self.model_sizes = [ + "tiny", "base", "small", "medium", "large-v2" + ] + 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 + device = "cuda" if torch.cuda.is_available() else "cpu" + + + if self.model_size == None: + return + self.transcriber = WhisperModel( - "small" if multilingual else "small.en", + self.model_size, device=device, compute_type="int8" if device=="cpu" else "float16", local_files_only=False, @@ -228,7 +251,6 @@ class ServeClient: self.pick_previous_segments = 2 # threading - self.websocket = websocket self.trans_thread = threading.Thread(target=self.speech_to_text) self.trans_thread.start() self.websocket.send( @@ -240,34 +262,30 @@ class ServeClient: ) ) - def fill_output(self, output): + def get_model_size(self, model_size): """ - Format the current incomplete transcription output by combining it with previous complete segments. - The resulting transcription is wrapped into two lines, each containing a maximum of 50 characters. - - It ensures that the combined transcription fits within two lines, with a maximum of 50 characters per line. - Segments are concatenated in the order they exist in the list of previous segments, with the most - recent complete segment first and older segments prepended as needed to maintain the character limit. - If a 3-second pause is detected in the previous segments, any text preceding it is discarded to ensure - the transcription starts with the most recent complete content. The resulting transcription is returned - as a single string. - - Args: - output(str): The current incomplete transcription segment. + Returns the whisper model size based on multilingual. + """ + 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 - Returns: - str: A formatted transcription wrapped in two lines. - """ - text = '' - pick_prev = min(len(self.text), self.pick_previous_segments) - for seg in self.text[-pick_prev:]: - # discard everything before a 3 second pause - if seg == '': - text = '' - else: - text += seg - wrapped = "".join(text + output) - return wrapped + if model_size == "large-v2": + self.multilingual = True + return model_size + + if not self.multilingual: + model_size = model_size + ".en" + + return model_size def add_frames(self, frame_np): """