create new method for handling a new connection; expcetion handling

This commit is contained in:
makaveli10
2024-02-15 12:16:18 +05:30
parent 170ba22e5b
commit 7bcb92c266
+46 -34
View File
@@ -7,7 +7,7 @@ import logging
import torch import torch
import numpy as np import numpy as np
from websockets.sync.server import serve from websockets.sync.server import serve
from websockets.exceptions import ConnectionClosed
from whisper_live.vad import VoiceActivityDetector from whisper_live.vad import VoiceActivityDetector
from whisper_live.transcriber import WhisperModel from whisper_live.transcriber import WhisperModel
try: try:
@@ -186,6 +186,23 @@ class TranscriptionServer:
frame_data = websocket.recv() frame_data = websocket.recv()
return np.frombuffer(frame_data, dtype=np.float32) return np.frombuffer(frame_data, dtype=np.float32)
def handle_new_connection(self, websocket, backend, faster_whisper_custom_model_path,
whisper_tensorrt_path, trt_multilingual):
logging.info("New client connected")
options = websocket.recv()
options = json.loads(options)
if self.client_manager.is_server_full(websocket, options):
websocket.close()
return
self.backend = backend
if self.backend == "tensorrt":
self.vad_detector = VoiceActivityDetector(frame_rate=self.RATE)
self.initialize_client(
websocket, options, faster_whisper_custom_model_path, whisper_tensorrt_path, trt_multilingual)
def recv_audio(self, def recv_audio(self,
websocket, websocket,
backend="faster_whisper", backend="faster_whisper",
@@ -216,45 +233,40 @@ class TranscriptionServer:
Raises: Raises:
Exception: If there is an error during the audio frame processing. Exception: If there is an error during the audio frame processing.
""" """
logging.info("New client connected") try:
options = websocket.recv() self.handle_new_connection(websocket, backend, faster_whisper_custom_model_path,
options = json.loads(options) whisper_tensorrt_path, trt_multilingual)
if self.client_manager.is_server_full(websocket, options): while not self.client_manager.is_client_timeout(websocket):
websocket.close() try:
return frame_np = self.get_audio_from_websocket(websocket)
client = self.client_manager.get_client(websocket)
self.backend = backend # VAD, for faster_whisper VAD model is already integrated
if self.backend == "tensorrt": if self.backend == "tensorrt":
self.vad_detector = VoiceActivityDetector(frame_rate=self.RATE) if not self.voice_activity(websocket, frame_np):
continue
self.no_voice_activity_chunks = 0
client.set_eos(False)
self.initialize_client( client.add_frames(frame_np)
websocket, options, faster_whisper_custom_model_path, whisper_tensorrt_path, trt_multilingual)
while not self.client_manager.is_client_timeout(websocket): except Exception as e:
try: logging.error(e)
frame_np = self.get_audio_from_websocket(websocket) self.cleanup(websocket)
client = self.client_manager.get_client(websocket) websocket.close()
break
# VAD, for faster_whisper VAD model is already integrated except ConnectionClosed:
if self.backend == "tensorrt": logging.info(f"Connection closed by client with path: {websocket.path}")
if not self.voice_activity(websocket, frame_np): except json.JSONDecodeError:
continue logging.error("Failed to decode JSON from client")
self.no_voice_activity_chunks = 0 except Exception as e:
client.set_eos(False) logging.error(f"Unexpected error: {str(e)}")
finally:
client.add_frames(frame_np) if self.client_manager.get_client(websocket):
except Exception as e:
logging.error(e)
self.cleanup(websocket) self.cleanup(websocket)
websocket.close() websocket.close()
break del websocket
if self.client_manager.get_client(websocket):
self.cleanup(websocket)
websocket.close()
del websocket
def run(self, def run(self,
host, host,