diff --git a/whisper_live/server.py b/whisper_live/server.py index 2accaac..b6efc77 100644 --- a/whisper_live/server.py +++ b/whisper_live/server.py @@ -4,6 +4,9 @@ import threading import json import functools import logging +from enum import Enum +from typing import List, Optional + import torch import numpy as np from websockets.sync.server import serve @@ -121,6 +124,25 @@ class ClientManager: return False +class BackendType(Enum): + FASTER_WHISPER = "faster_whisper" + TENSORRT = "tensorrt" + + @staticmethod + def valid_types() -> List[str]: + return [backend_type.value for backend_type in BackendType] + + @staticmethod + def is_valid(backend: str) -> bool: + return backend in BackendType.valid_types() + + def is_faster_whisper(self) -> bool: + return self == BackendType.FASTER_WHISPER + + def is_tensorrt(self) -> bool: + return self == BackendType.TENSORRT + + class TranscriptionServer: RATE = 16000 @@ -134,7 +156,9 @@ class TranscriptionServer: self, websocket, options, faster_whisper_custom_model_path, whisper_tensorrt_path, trt_multilingual ): - if self.backend == "tensorrt": + client: Optional[ServeClientBase] = None + + if self.backend.is_tensorrt(): try: client = ServeClientTensorRT( websocket, @@ -155,9 +179,9 @@ class TranscriptionServer: "message": "TensorRT-LLM not supported on Server yet. " "Reverting to available backend: 'faster_whisper'" })) - self.backend = "faster_whisper" + self.backend = BackendType.FASTER_WHISPER - if self.backend == "faster_whisper": + if self.backend.is_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 @@ -174,6 +198,9 @@ class TranscriptionServer: ) logging.info("Running faster_whisper backend.") + if client is None: + raise ValueError(f"Backend type {self.backend.value} not recognised or not handled.") + self.client_manager.add_client(websocket, client) def get_audio_from_websocket(self, websocket): @@ -202,7 +229,7 @@ class TranscriptionServer: websocket.close() return False # Indicates that the connection should not continue - if self.backend == "tensorrt": + if self.backend.is_tensorrt(): self.vad_detector = VoiceActivityDetector(frame_rate=self.RATE) self.initialize_client(websocket, options, faster_whisper_custom_model_path, whisper_tensorrt_path, trt_multilingual) @@ -221,11 +248,11 @@ class TranscriptionServer: frame_np = self.get_audio_from_websocket(websocket) client = self.client_manager.get_client(websocket) if frame_np is False: - if self.backend == "tensorrt": + if self.backend.is_tensorrt(): client.set_eos(True) return False - if self.backend == "tensorrt": + if self.backend.is_tensorrt(): voice_active = self.voice_activity(websocket, frame_np) if voice_active: self.no_voice_activity_chunks = 0 @@ -238,7 +265,7 @@ class TranscriptionServer: def recv_audio(self, websocket, - backend="faster_whisper", + backend: BackendType = BackendType.FASTER_WHISPER, faster_whisper_custom_model_path=None, whisper_tensorrt_path=None, trt_multilingual=False): @@ -311,10 +338,12 @@ class TranscriptionServer: # TODO: load model initially else: logging.info("Single model mode currently only works with custom models.") + if not BackendType.is_valid(backend): + raise ValueError(f"{backend} is not a valid backend type. Choose backend from {BackendType.valid_types()}") with serve( functools.partial( self.recv_audio, - backend=backend, + backend=BackendType(backend), faster_whisper_custom_model_path=faster_whisper_custom_model_path, whisper_tensorrt_path=whisper_tensorrt_path, trt_multilingual=trt_multilingual