add: tensorrt backend to server

This commit is contained in:
makaveli10
2024-01-09 18:10:17 +00:00
parent fd86340f30
commit 0f9e93d203
+20 -2
View File
@@ -1,5 +1,23 @@
import argparse
from whisper_live.server import TranscriptionServer
if __name__ == "__main__":
server = TranscriptionServer()
server.run("0.0.0.0")
parser = argparse.ArgumentParser()
parser.add_argument('--backend', type=str, default='tensorrt', help='Backends from ["tensorrt", "faster_whisper"]')
parser.add_argument('--whisper_tensorrt_path',
type=str,
default="/root/TensorRT-LLM/examples/whisper/whisper_small_en",
help='Whisper TensorRT model path')
args = parser.parse_args()
if args.backend == "tensorrt":
if args.whisper_tensorrt_path is None:
raise ValueError("Please Provide a valid tensorrt model path")
from whisper_live.server_tensorrt import TranscriptionServerTRT
server = TranscriptionServerTRT()
server.run("0.0.0.0", port=6006, whisper_tensorrt_path=args.whisper_tensorrt_path)
else:
server = TranscriptionServer()
server.run("0.0.0.0")