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
+19 -9
View File
@@ -3,26 +3,36 @@ from whisper_live.server import TranscriptionServer
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--port', type=int, default=9090, help="Websocket port to run the server on.")
parser.add_argument('--backend', type=str, default='faster_whisper', help='Backends from ["tensorrt", "faster_whisper"]')
parser.add_argument('--whisper_tensorrt_path',
parser.add_argument('--port', '-p',
type=int,
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,
default=None,
help='Whisper TensorRT model path')
parser.add_argument('--trt_multilingual',
parser.add_argument('--trt_multilingual', '-m',
action="store_true",
help='Boolean only for TensorRT model. True if multilingual.')
args = parser.parse_args()
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")
server = TranscriptionServer()
server.run(
"0.0.0.0",
port=6006,
backend=args.backend,
whisper_tensorrt_path=args.whisper_tensorrt_path,
multilingual=args.trt_multilingual
port=args.port,
backend=args.backend,
faster_whisper_custom_model_path=args.faster_whisper_custom_model_path,
whisper_tensorrt_path=args.trt_model_path,
trt_multilingual=args.trt_multilingual
)
whisper_live/server.py
+46 -25
View File
@@ -1,3 +1,4 @@
import os
import websockets
import time
import threading
@@ -68,7 +69,12 @@ class TranscriptionServer:
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.
@@ -85,7 +91,11 @@ class TranscriptionServer:
Args:
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:
Exception: If there is an error during the audio frame processing.
"""
@@ -118,11 +128,11 @@ class TranscriptionServer:
self.backend = "tensorrt"
client = ServeClientTensorRT(
websocket,
multilingual=multilingual,
multilingual=trt_multilingual,
language=options["language"],
task=options["task"],
client_uid=options["uid"],
model_path=whisper_tensorrt_path
model=whisper_tensorrt_path
)
logging.info(f"Running TensorRT backend.")
except Exception as e:
@@ -138,13 +148,17 @@ class TranscriptionServer:
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(
websocket,
multilingual=options["multilingual"],
language=options["language"],
task=options["task"],
client_uid=options["uid"],
model_size=options["model_size"],
model=options["model"],
initial_prompt=options.get("initial_prompt"),
vad_parameters=options.get("vad_parameters")
)
@@ -159,7 +173,7 @@ class TranscriptionServer:
frame_data = websocket.recv()
frame_np = np.frombuffer(frame_data, dtype=np.float32)
# VAD
# VAD, for faster_whisper VAD model is already integrated
if self.backend == "tensorrt":
try:
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 not self.clients[websocket].eos:
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
no_voice_activity_chunks = 0
self.clients[websocket].set_eos(False)
@@ -198,7 +212,14 @@ class TranscriptionServer:
del websocket
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.
@@ -210,8 +231,9 @@ class TranscriptionServer:
functools.partial(
self.recv_audio,
backend=backend,
faster_whisper_custom_model_path=faster_whisper_custom_model_path,
whisper_tensorrt_path=whisper_tensorrt_path,
multilingual=multilingual
trt_multilingual=trt_multilingual
),
host,
port
@@ -220,6 +242,10 @@ class TranscriptionServer:
class ServeClientBase(object):
RATE = 16000
SERVER_READY = "SERVER_READY"
DISCONNECT = "DISCONNECT"
def __init__(self, client_uid, websocket):
self.client_uid = client_uid
self.websocket = websocket
@@ -303,7 +329,6 @@ class ServeClientBase(object):
"""
logging.info("Cleaning up.")
self.exit = True
self.transcriber.destroy()
class ServeClientTensorRT(ServeClientBase):
@@ -335,10 +360,6 @@ class ServeClientTensorRT(ServeClientBase):
pick_previous_segments (int): Number of previous segments to include in the output.
websocket: The WebSocket connection for the client.
"""
RATE = 16000
SERVER_READY = "SERVER_READY"
DISCONNECT = "DISCONNECT"
def __init__(
self,
websocket,
@@ -347,7 +368,7 @@ class ServeClientTensorRT(ServeClientBase):
multilingual=False,
language=None,
client_uid=None,
model_path=None
model=None
):
"""
Initialize a ServeClient instance.
@@ -369,7 +390,7 @@ class ServeClientTensorRT(ServeClientBase):
self.task = task
self.eos = False
self.transcriber = WhisperTRTLLM(
model_path,
model,
assets_dir="assets",
device="cuda",
is_multilingual=multilingual,
@@ -530,10 +551,6 @@ class ServeClientFasterWhisper(ServeClientBase):
pick_previous_segments (int): Number of previous segments to include in the output.
websocket: The WebSocket connection for the client.
"""
RATE = 16000
SERVER_READY = "SERVER_READY"
DISCONNECT = "DISCONNECT"
def __init__(
self,
websocket,
@@ -542,7 +559,7 @@ class ServeClientFasterWhisper(ServeClientBase):
multilingual=False,
language=None,
client_uid=None,
model_size="small",
model="small",
initial_prompt=None,
vad_parameters=None
):
@@ -563,10 +580,14 @@ class ServeClientFasterWhisper(ServeClientBase):
"""
super().__init__(client_uid, websocket)
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.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.task = task
self.initial_prompt = initial_prompt
@@ -574,11 +595,11 @@ class ServeClientFasterWhisper(ServeClientBase):
device = "cuda" if torch.cuda.is_available() else "cpu"
if self.model_size == None:
if self.model_size_or_path == None:
return
self.transcriber = WhisperModel(
self.model_size,
self.model_size_or_path,
device=device,
compute_type="int8" if device=="cpu" else "float16",
local_files_only=False,