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

merge with main

Browse Source
This commit is contained in:
makaveli10
2024-01-19 11:43:21 +00:00
parent 67232fffd5
commit 735d6c7763
2 changed files with 65 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
run_server.py
+18 -8
View File
@@ -3,26 +3,36 @@ from whisper_live.server import TranscriptionServer
if __name__ == "__main__": if __name__ == "__main__":
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--port', type=int, default=9090, help="Websocket port to run the server on.") parser.add_argument('--port', '-p',
parser.add_argument('--backend', type=str, default='faster_whisper', help='Backends from ["tensorrt", "faster_whisper"]') type=int,
parser.add_argument('--whisper_tensorrt_path', 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, type=str,
default=None, default=None,
help='Whisper TensorRT model path') help='Whisper TensorRT model path')
parser.add_argument('--trt_multilingual', parser.add_argument('--trt_multilingual', '-m',
action="store_true", action="store_true",
help='Boolean only for TensorRT model. True if multilingual.') help='Boolean only for TensorRT model. True if multilingual.')
args = parser.parse_args() args = parser.parse_args()
if args.backend == "tensorrt": 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") raise ValueError("Please Provide a valid tensorrt model path")
server = TranscriptionServer() server = TranscriptionServer()
server.run( server.run(
"0.0.0.0", "0.0.0.0",
port=6006, port=args.port,
backend=args.backend, backend=args.backend,
whisper_tensorrt_path=args.whisper_tensorrt_path, faster_whisper_custom_model_path=args.faster_whisper_custom_model_path,
multilingual=args.trt_multilingual whisper_tensorrt_path=args.trt_model_path,
trt_multilingual=args.trt_multilingual
) )
whisper_live/server.py
+45 -24
View File
@@ -1,3 +1,4 @@
import os
import websockets import websockets
import time import time
import threading import threading
@@ -68,7 +69,12 @@ class TranscriptionServer:
return wait_time / 60 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. Receive audio chunks from a client in an infinite loop.
@@ -85,6 +91,10 @@ class TranscriptionServer:
Args: Args:
websocket (WebSocket): The WebSocket connection for the client. 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: Raises:
Exception: If there is an error during the audio frame processing. Exception: If there is an error during the audio frame processing.
@@ -118,11 +128,11 @@ class TranscriptionServer:
self.backend = "tensorrt" self.backend = "tensorrt"
client = ServeClientTensorRT( client = ServeClientTensorRT(
websocket, websocket,
multilingual=multilingual, multilingual=trt_multilingual,
language=options["language"], language=options["language"],
task=options["task"], task=options["task"],
client_uid=options["uid"], client_uid=options["uid"],
model_path=whisper_tensorrt_path model=whisper_tensorrt_path
) )
logging.info(f"Running TensorRT backend.") logging.info(f"Running TensorRT backend.")
except Exception as e: except Exception as e:
@@ -138,13 +148,17 @@ class TranscriptionServer:
self.backend = "faster_whisper" self.backend = "faster_whisper"
if 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( client = ServeClientFasterWhisper(
websocket, websocket,
multilingual=options["multilingual"], multilingual=options["multilingual"],
language=options["language"], language=options["language"],
task=options["task"], task=options["task"],
client_uid=options["uid"], client_uid=options["uid"],
model_size=options["model_size"], model=options["model"],
initial_prompt=options.get("initial_prompt"), initial_prompt=options.get("initial_prompt"),
vad_parameters=options.get("vad_parameters") vad_parameters=options.get("vad_parameters")
) )
@@ -159,7 +173,7 @@ class TranscriptionServer:
frame_data = websocket.recv() frame_data = websocket.recv()
frame_np = np.frombuffer(frame_data, dtype=np.float32) frame_np = np.frombuffer(frame_data, dtype=np.float32)
# VAD # VAD, for faster_whisper VAD model is already integrated
if self.backend == "tensorrt": if self.backend == "tensorrt":
try: try:
speech_prob = self.vad_model(torch.from_numpy(frame_np.copy()), self.RATE).item() 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 no_voice_activity_chunks > 3:
if not self.clients[websocket].eos: if not self.clients[websocket].eos:
self.clients[websocket].set_eos(True) 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 continue
no_voice_activity_chunks = 0 no_voice_activity_chunks = 0
self.clients[websocket].set_eos(False) self.clients[websocket].set_eos(False)
@@ -198,7 +212,14 @@ class TranscriptionServer:
del websocket del websocket
break 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. Run the transcription server.
@@ -210,8 +231,9 @@ class TranscriptionServer:
functools.partial( functools.partial(
self.recv_audio, self.recv_audio,
backend=backend, backend=backend,
faster_whisper_custom_model_path=faster_whisper_custom_model_path,
whisper_tensorrt_path=whisper_tensorrt_path, whisper_tensorrt_path=whisper_tensorrt_path,
multilingual=multilingual trt_multilingual=trt_multilingual
), ),
host, host,
port port
@@ -220,6 +242,10 @@ class TranscriptionServer:
class ServeClientBase(object): class ServeClientBase(object):
RATE = 16000
SERVER_READY = "SERVER_READY"
DISCONNECT = "DISCONNECT"
def __init__(self, client_uid, websocket): def __init__(self, client_uid, websocket):
self.client_uid = client_uid self.client_uid = client_uid
self.websocket = websocket self.websocket = websocket
@@ -303,7 +329,6 @@ class ServeClientBase(object):
""" """
logging.info("Cleaning up.") logging.info("Cleaning up.")
self.exit = True self.exit = True
self.transcriber.destroy()
class ServeClientTensorRT(ServeClientBase): class ServeClientTensorRT(ServeClientBase):
@@ -335,10 +360,6 @@ class ServeClientTensorRT(ServeClientBase):
pick_previous_segments (int): Number of previous segments to include in the output. pick_previous_segments (int): Number of previous segments to include in the output.
websocket: The WebSocket connection for the client. websocket: The WebSocket connection for the client.
""" """
RATE = 16000
SERVER_READY = "SERVER_READY"
DISCONNECT = "DISCONNECT"
def __init__( def __init__(
self, self,
websocket, websocket,
@@ -347,7 +368,7 @@ class ServeClientTensorRT(ServeClientBase):
multilingual=False, multilingual=False,
language=None, language=None,
client_uid=None, client_uid=None,
model_path=None model=None
): ):
""" """
Initialize a ServeClient instance. Initialize a ServeClient instance.
@@ -369,7 +390,7 @@ class ServeClientTensorRT(ServeClientBase):
self.task = task self.task = task
self.eos = False self.eos = False
self.transcriber = WhisperTRTLLM( self.transcriber = WhisperTRTLLM(
model_path, model,
assets_dir="assets", assets_dir="assets",
device="cuda", device="cuda",
is_multilingual=multilingual, is_multilingual=multilingual,
@@ -530,10 +551,6 @@ class ServeClientFasterWhisper(ServeClientBase):
pick_previous_segments (int): Number of previous segments to include in the output. pick_previous_segments (int): Number of previous segments to include in the output.
websocket: The WebSocket connection for the client. websocket: The WebSocket connection for the client.
""" """
RATE = 16000
SERVER_READY = "SERVER_READY"
DISCONNECT = "DISCONNECT"
def __init__( def __init__(
self, self,
websocket, websocket,
@@ -542,7 +559,7 @@ class ServeClientFasterWhisper(ServeClientBase):
multilingual=False, multilingual=False,
language=None, language=None,
client_uid=None, client_uid=None,
model_size="small", model="small",
initial_prompt=None, initial_prompt=None,
vad_parameters=None vad_parameters=None
): ):
@@ -563,10 +580,14 @@ class ServeClientFasterWhisper(ServeClientBase):
""" """
super().__init__(client_uid, websocket) super().__init__(client_uid, websocket)
self.model_sizes = [ 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.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.language = language if self.multilingual else "en"
self.task = task self.task = task
self.initial_prompt = initial_prompt self.initial_prompt = initial_prompt
@@ -574,11 +595,11 @@ class ServeClientFasterWhisper(ServeClientBase):
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"
if self.model_size == None: if self.model_size_or_path == None:
return return
self.transcriber = WhisperModel( self.transcriber = WhisperModel(
self.model_size, self.model_size_or_path,
device=device, device=device,
compute_type="int8" if device=="cpu" else "float16", compute_type="int8" if device=="cpu" else "float16",
local_files_only=False, local_files_only=False,