cjones/WhisperLive
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

update client to handle server messages

Browse Source
This commit is contained in:
makaveli10
2023-08-09 01:35:23 +05:30
parent defc1687b7
commit e64e26fe47
1 changed files with 30 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
whisper_live/client.py
+30 -6
View File
@@ -10,6 +10,7 @@ import threading
import textwrap import textwrap
import json import json
import websocket import websocket
import uuid
def resample(file: str, sr: int = 16000): def resample(file: str, sr: int = 16000):
@@ -47,10 +48,12 @@ class Client:
CHANNELS = 1 CHANNELS = 1
RATE = 16000 RATE = 16000
RECORD_SECONDS = 60000 RECORD_SECONDS = 60000
START_RECORDING = False RECORDING = False
multilingual = False multilingual = False
language = None language = None
task = "transcribe" task = "transcribe"
uid = str(uuid.uuid4())
WAITING = False
def __init__(self, host=None, port=None, is_multilingual=False, lang=None, translate=False): def __init__(self, host=None, port=None, is_multilingual=False, lang=None, translate=False):
Client.multilingual = is_multilingual Client.multilingual = is_multilingual
@@ -90,16 +93,32 @@ class Client:
@staticmethod @staticmethod
def on_message(ws, message): def on_message(ws, message):
message = json.loads(message) message = json.loads(message)
if message == "SERVER_READY": if message.get('uid')!=Client.uid:
Client.START_RECORDING = True print("[ERROR]: invalid client uid")
return return
if isinstance(message, dict): if "status" in message.keys() and message["status"] == "WAIT":
Client.WAITING = True
print(f"[INFO]:Server is full. Estimated wait time {round(message['message'])} minutes.")
if "message" in message.keys() and message["message"] == "DISCONNECT":
print("[INFO]: Server overtime disconnected.")
Client.RECORDING = False
if "message" in message.keys() and message["message"] == "SERVER_READY":
Client.RECORDING = True
return
if "language" in message.keys():
Client.language = message.get("language") Client.language = message.get("language")
lang_prob = message.get("language_prob") lang_prob = message.get("language_prob")
print(f"[INFO]: Server detected language {Client.language} with probability {lang_prob}") print(f"[INFO]: Server detected language {Client.language} with probability {lang_prob}")
return return
if "segments" not in message.keys():
return
message = message["segments"]
text = [] text = []
if len(message): if len(message):
for seg in message: for seg in message:
@@ -134,6 +153,7 @@ class Client:
print("[INFO]: Opened connection") print("[INFO]: Opened connection")
ws.send(json.dumps({ ws.send(json.dumps({
'uid': Client.uid,
'multilingual': Client.multilingual, 'multilingual': Client.multilingual,
'language': Client.language, 'language': Client.language,
'task': Client.task 'task': Client.task
@@ -162,7 +182,7 @@ class Client:
output=True, output=True,
frames_per_buffer=self.CHUNK) frames_per_buffer=self.CHUNK)
try: try:
while True: while Client.RECORDING:
data = self.wf.readframes(self.CHUNK) data = self.wf.readframes(self.CHUNK)
if data==b'': break if data==b'': break
@@ -210,6 +230,7 @@ class Client:
os.makedirs("chunks", exist_ok=True) os.makedirs("chunks", exist_ok=True)
try: try:
for _ in range(0, int(self.RATE / self.CHUNK * self.RECORD_SECONDS)): for _ in range(0, int(self.RATE / self.CHUNK * self.RECORD_SECONDS)):
if not Client.RECORDING: break
data = self.stream.read(self.CHUNK) data = self.stream.read(self.CHUNK)
self.frames += data self.frames += data
@@ -264,7 +285,10 @@ class TranscriptionClient:
def __call__(self, audio=None): def __call__(self, audio=None):
print("[INFO]: Waiting for server ready ...") print("[INFO]: Waiting for server ready ...")
while not Client.START_RECORDING: while not Client.RECORDING:
if Client.WAITING:
self.client.close_websocket()
return
pass pass
print("[INFO]: Server Ready!") print("[INFO]: Server Ready!")
if audio is not None: if audio is not None: