From 47ee035f65cd74ca42821f89b8621781cf935158 Mon Sep 17 00:00:00 2001 From: makaveli10 Date: Tue, 22 Apr 2025 12:33:03 +0000 Subject: [PATCH 1/2] Upgrade tensorrt_llm to v0.18.2 Signed-off-by: makaveli10 --- README.md | 9 +- TensorRT_whisper.md | 10 +- docker/Dockerfile.tensorrt | 15 ++- run_server.py | 4 + scripts/build_whisper_tensorrt.sh | 8 +- whisper_live/backend/trt_backend.py | 28 ++++-- whisper_live/server.py | 18 ++-- .../transcriber/transcriber_tensorrt.py | 99 ++++++++++++++----- 8 files changed, 134 insertions(+), 57 deletions(-) diff --git a/README.md b/README.md index 2119540..8b4389c 100644 --- a/README.md +++ b/README.md @@ -141,16 +141,11 @@ client(hls_url="http://as-hls-ww-live.akamaized.net/pool_904/live/ww/bbc_1xtra/b ## Browser Extensions - Run the server with your desired backend as shown [here](https://github.com/collabora/WhisperLive?tab=readme-ov-file#running-the-server). -- Transcribe audio directly from your browser using our Chrome or Firefox extensions. Refer to [Audio-Transcription-Chrome](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Chrome#readme) and [Audio-Transcription-Firefox](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Firefox#readme) for setup instructions. - -## Whisper Live Server in Docker -- GPU - - Faster-Whisper - ```bash +- Transcribe audio directly from your browser using our Chrome or Firefox extensions. Refer to [Audio-Transcription-Chrome](https://github.com/collabora/whisper-live/tree/main/Audio-Transcription-Chrome#readme) and https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md docker run -it --gpus all -p 9090:9090 ghcr.io/collabora/whisperlive-gpu:latest ``` - - TensorRT. + - TensorRT. Refer to [TensorRT_whisper readme](https://github.com/collabora/WhisperLive/blob/main/TensorRT_whisper.md) for setup and more tensorrt backend configurations. ```bash docker run -p 9090:9090 --runtime=nvidia --gpus all --entrypoint /bin/bash -it ghcr.io/collabora/whisperlive-tensorrt diff --git a/TensorRT_whisper.md b/TensorRT_whisper.md index 56a0318..785f1f4 100644 --- a/TensorRT_whisper.md +++ b/TensorRT_whisper.md @@ -1,6 +1,6 @@ # WhisperLive-TensorRT We have only tested the TensorRT backend in docker so, we recommend docker for a smooth TensorRT backend setup. -**Note**: We use `tensorrt_llm==0.15.0.dev2024111200` +**Note**: We use `tensorrt_llm==0.18.2` ## Installation - Install [docker](https://docs.docker.com/engine/install/) @@ -36,3 +36,11 @@ python3 run_server.py --port 9090 \ --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_float16" \ --trt_multilingual ``` + +By default trt_backend uses cpp_session, to use python session pass `--trt_py_session` to run_server.py +```bash +python3 run_server.py --port 9090 \ + --backend tensorrt \ + --trt_model_path "/app/TensorRT-LLM-examples/whisper/whisper_small_float16" \ + --trt_py_session +``` \ No newline at end of file diff --git a/docker/Dockerfile.tensorrt b/docker/Dockerfile.tensorrt index c46a2d2..3b7c880 100644 --- a/docker/Dockerfile.tensorrt +++ b/docker/Dockerfile.tensorrt @@ -1,19 +1,19 @@ -FROM nvidia/cuda:12.4.1-base-ubuntu22.04 AS base +FROM nvidia/cuda:12.8.1-base-ubuntu22.04 AS base ARG DEBIAN_FRONTEND=noninteractive RUN apt-get update && apt-get install -y \ python3.10 python3-pip openmpi-bin libopenmpi-dev git git-lfs wget \ + && apt install python-is-python3 \ + && pip install --upgrade pip setuptools \ && rm -rf /var/lib/apt/lists/* FROM base AS devel -RUN pip3 install --no-cache-dir -U tensorrt_llm==0.15.0.dev2024111200 --extra-index-url https://pypi.nvidia.com +RUN pip install --no-cache-dir -U tensorrt_llm==0.18.2 --extra-index-url https://pypi.nvidia.com WORKDIR /app -RUN git clone https://github.com/NVIDIA/TensorRT-LLM.git && cd TensorRT-LLM && \ - git checkout c629546ce429623c8a163633095230154a6f0574 && cd ../ && \ - mv TensorRT-LLM/examples ./TensorRT-LLM-examples && \ - rm -rf TensorRT-LLM - +RUN git clone -b v0.18.2 https://github.com/NVIDIA/TensorRT-LLM.git \ + && mv TensorRT-LLM/examples ./TensorRT-LLM-examples \ + && rm -rf TensorRT-LLM FROM devel AS release WORKDIR /app @@ -25,7 +25,6 @@ RUN apt update && bash setup.sh && rm setup.sh COPY requirements/server.txt . RUN pip install --no-cache-dir -r server.txt && rm server.txt -RUN pip install pynvml==11.5.0 COPY whisper_live ./whisper_live COPY scripts/build_whisper_tensorrt.sh . COPY run_server.py . \ No newline at end of file diff --git a/run_server.py b/run_server.py index 2dd5581..db093ef 100644 --- a/run_server.py +++ b/run_server.py @@ -21,6 +21,9 @@ if __name__ == "__main__": parser.add_argument('--trt_multilingual', '-m', action="store_true", help='Boolean only for TensorRT model. True if multilingual.') + parser.add_argument('--trt_py_session', + action="store_true", + help='Boolean only for TensorRT model. Use python session or cpp session, By default uses Cpp.') parser.add_argument('--omp_num_threads', '-omp', type=int, default=1, @@ -46,5 +49,6 @@ if __name__ == "__main__": faster_whisper_custom_model_path=args.faster_whisper_custom_model_path, whisper_tensorrt_path=args.trt_model_path, trt_multilingual=args.trt_multilingual, + trt_py_session=args.trt_py_session, single_model=not args.no_single_model, ) diff --git a/scripts/build_whisper_tensorrt.sh b/scripts/build_whisper_tensorrt.sh index 0b5c22c..e09e871 100644 --- a/scripts/build_whisper_tensorrt.sh +++ b/scripts/build_whisper_tensorrt.sh @@ -54,7 +54,7 @@ download_and_build_model() { local inference_precision="float16" local weight_only_precision="${2:-float16}" local max_beam_width=4 - local max_batch_size=1 + local max_batch_size=4 echo "Downloading $model_name..." # wget --directory-prefix=assets "$model_url" @@ -80,7 +80,6 @@ download_and_build_model() { --checkpoint_dir "${checkpoint_dir}/encoder" \ --output_dir "${output_dir}/encoder" \ --moe_plugin disable \ - --enable_xqa disable \ --max_batch_size "$max_batch_size" \ --gemm_plugin disable \ --bert_attention_plugin "$inference_precision" \ @@ -92,11 +91,10 @@ download_and_build_model() { --checkpoint_dir "${checkpoint_dir}/decoder" \ --output_dir "${output_dir}/decoder" \ --moe_plugin disable \ - --enable_xqa disable \ --max_beam_width "$max_beam_width" \ --max_batch_size "$max_batch_size" \ - --max_seq_len 200 \ - --max_input_len 14 \ + --max_seq_len 225 \ + --max_input_len 32 \ --max_encoder_input_len 3000 \ --gemm_plugin "$inference_precision" \ --bert_attention_plugin "$inference_precision" \ diff --git a/whisper_live/backend/trt_backend.py b/whisper_live/backend/trt_backend.py index dc7665a..a24e2a9 100644 --- a/whisper_live/backend/trt_backend.py +++ b/whisper_live/backend/trt_backend.py @@ -11,7 +11,18 @@ class ServeClientTensorRT(ServeClientBase): SINGLE_MODEL = None SINGLE_MODEL_LOCK = threading.Lock() - def __init__(self, websocket, task="transcribe", multilingual=False, language=None, client_uid=None, model=None, single_model=False): + def __init__( + self, + websocket, + task="transcribe", + multilingual=False, + language=None, + client_uid=None, + model=None, + single_model=False, + use_py_session=False, + max_new_tokens=225, + ): """ Initialize a ServeClient instance. The Whisper model is initialized based on the client's language and device availability. @@ -26,21 +37,24 @@ class ServeClientTensorRT(ServeClientBase): language (str, optional): The language for transcription. Defaults to None. client_uid (str, optional): A unique identifier for the client. Defaults to None. single_model (bool, optional): Whether to instantiate a new model for each client connection. Defaults to False. + use_py_session (bool, optional): Use python session or cpp session. Defaults to Cpp Session. + max_new_tokens (int, optional): Max number of tokens to generate. """ super().__init__(client_uid, websocket) self.language = language if multilingual else "en" self.task = task self.eos = False + self.max_new_tokens = max_new_tokens if single_model: if ServeClientTensorRT.SINGLE_MODEL is None: - self.create_model(model, multilingual) + self.create_model(model, multilingual, use_py_session=use_py_session) ServeClientTensorRT.SINGLE_MODEL = self.transcriber else: self.transcriber = ServeClientTensorRT.SINGLE_MODEL else: - self.create_model(model, multilingual) + self.create_model(model, multilingual, use_py_session=use_py_session) # threading self.trans_thread = threading.Thread(target=self.speech_to_text) @@ -52,7 +66,7 @@ class ServeClientTensorRT(ServeClientBase): "backend": "tensorrt" })) - def create_model(self, model, multilingual, warmup=True): + def create_model(self, model, multilingual, warmup=True, use_py_session=False): """ Instantiates a new model, sets it as the transcriber and does warmup if desired. """ @@ -62,7 +76,9 @@ class ServeClientTensorRT(ServeClientBase): device="cuda", is_multilingual=multilingual, language=self.language, - task=self.task + task=self.task, + use_py_session=use_py_session, + max_output_len=self.max_new_tokens, ) if warmup: self.warmup() @@ -117,7 +133,7 @@ class ServeClientTensorRT(ServeClientBase): mel, duration = self.transcriber.log_mel_spectrogram(input_bytes) last_segment = self.transcriber.transcribe( mel, - text_prefix=f"<|startoftranscript|><|{self.language}|><|{self.task}|><|notimestamps|>" + text_prefix=f"<|startoftranscript|><|{self.language}|><|{self.task}|><|notimestamps|>", ) if ServeClientTensorRT.SINGLE_MODEL: ServeClientTensorRT.SINGLE_MODEL_LOCK.release() diff --git a/whisper_live/server.py b/whisper_live/server.py index 5a59ff7..e0fee07 100644 --- a/whisper_live/server.py +++ b/whisper_live/server.py @@ -153,7 +153,7 @@ class TranscriptionServer: def initialize_client( self, websocket, options, faster_whisper_custom_model_path, - whisper_tensorrt_path, trt_multilingual + whisper_tensorrt_path, trt_multilingual, trt_py_session=False, ): client: Optional[ServeClientBase] = None @@ -168,6 +168,7 @@ class TranscriptionServer: client_uid=options["uid"], model=whisper_tensorrt_path, single_model=self.single_model, + use_py_session=trt_py_session, ) logging.info("Running TensorRT backend.") except Exception as e: @@ -248,7 +249,7 @@ class TranscriptionServer: return np.frombuffer(frame_data, dtype=np.float32) def handle_new_connection(self, websocket, faster_whisper_custom_model_path, - whisper_tensorrt_path, trt_multilingual): + whisper_tensorrt_path, trt_multilingual, trt_py_session=False): try: logging.info("New client connected") options = websocket.recv() @@ -267,7 +268,7 @@ class TranscriptionServer: if self.backend.is_tensorrt(): self.vad_detector = VoiceActivityDetector(frame_rate=self.RATE) self.initialize_client(websocket, options, faster_whisper_custom_model_path, - whisper_tensorrt_path, trt_multilingual) + whisper_tensorrt_path, trt_multilingual, trt_py_session=trt_py_session) return True except json.JSONDecodeError: logging.error("Failed to decode JSON from client") @@ -299,11 +300,12 @@ class TranscriptionServer: return True def recv_audio(self, - websocket, + websocket, backend: BackendType = BackendType.FASTER_WHISPER, faster_whisper_custom_model_path=None, whisper_tensorrt_path=None, - trt_multilingual=False): + trt_multilingual=False, + trt_py_session=False): """ Receive audio chunks from a client in an infinite loop. @@ -330,7 +332,7 @@ class TranscriptionServer: """ self.backend = backend if not self.handle_new_connection(websocket, faster_whisper_custom_model_path, - whisper_tensorrt_path, trt_multilingual): + whisper_tensorrt_path, trt_multilingual, trt_py_session=trt_py_session): return try: @@ -354,6 +356,7 @@ class TranscriptionServer: faster_whisper_custom_model_path=None, whisper_tensorrt_path=None, trt_multilingual=False, + trt_py_session=False, single_model=False): """ Run the transcription server. @@ -381,7 +384,8 @@ class TranscriptionServer: backend=BackendType(backend), faster_whisper_custom_model_path=faster_whisper_custom_model_path, whisper_tensorrt_path=whisper_tensorrt_path, - trt_multilingual=trt_multilingual + trt_multilingual=trt_multilingual, + trt_py_session=trt_py_session, ), host, port diff --git a/whisper_live/transcriber/transcriber_tensorrt.py b/whisper_live/transcriber/transcriber_tensorrt.py index 7123534..fd5cb8f 100644 --- a/whisper_live/transcriber/transcriber_tensorrt.py +++ b/whisper_live/transcriber/transcriber_tensorrt.py @@ -23,7 +23,8 @@ from tensorrt_llm._utils import (str_dtype_to_torch, str_dtype_to_trt, from tensorrt_llm.bindings import GptJsonConfig, KVCacheType from tensorrt_llm.runtime import PYTHON_BINDINGS, ModelConfig, SamplingConfig from tensorrt_llm.runtime.session import Session, TensorInfo - +if PYTHON_BINDINGS: + from tensorrt_llm.runtime import ModelRunnerCpp SAMPLE_RATE = 16000 N_FFT = 400 @@ -255,8 +256,17 @@ class WhisperDecoding: class WhisperTRTLLM(object): - def __init__(self, engine_dir, assets_dir=None, device=None, is_multilingual=False, - language="en", task="transcribe"): + def __init__(self, + engine_dir, + assets_dir=None, + device=None, + is_multilingual=False, + language="en", + task="transcribe", + use_py_session=False, + num_beams=1, + debug_mode=False, + max_output_len=96): world_size = 1 runtime_rank = tensorrt_llm.mpi_rank() runtime_mapping = tensorrt_llm.Mapping(world_size, runtime_rank) @@ -268,13 +278,6 @@ class WhisperTRTLLM(object): self.num_languages = encoder_config['num_languages'] is_multilingual = (decoder_config['vocab_size'] >= 51865) - self.encoder = WhisperEncoding(engine_dir) - self.decoder = WhisperDecoding(engine_dir, - runtime_mapping, - debug_mode=False) - self.n_mels = self.encoder.n_mels - # self.tokenizer = get_tokenizer(num_languages=self.encoder.num_languages, - # tokenizer_dir=assets_dir) self.device = device self.tokenizer = get_tokenizer( is_multilingual, @@ -282,7 +285,28 @@ class WhisperTRTLLM(object): language=language, task=task, ) - self.filters = mel_filters(self.device, self.encoder.n_mels, assets_dir) + + if use_py_session: + self.encoder = WhisperEncoding(engine_dir) + self.decoder = WhisperDecoding(engine_dir, + runtime_mapping, + debug_mode=False) + else: + json_config = GptJsonConfig.parse_file(engine_dir / 'decoder' / + 'config.json') + assert json_config.model_config.supports_inflight_batching + runner_kwargs = dict(engine_dir=engine_dir, + is_enc_dec=True, + max_batch_size=1, + max_input_len=3000, + max_output_len=max_output_len, + max_beam_width=num_beams, + debug_mode=debug_mode, + kv_cache_free_gpu_memory_fraction=0.9, + cross_kv_cache_fraction=0.5) + self.model_runner_cpp = ModelRunnerCpp.from_dir(**runner_kwargs) + self.filters = mel_filters(self.device, self.n_mels, assets_dir) + self.use_py_session = use_py_session def log_mel_spectrogram( self, @@ -355,16 +379,38 @@ class WhisperTRTLLM(object): prompt_id = torch.tensor(prompt_id) batch_size = mel.shape[0] decoder_input_ids = prompt_id.repeat(batch_size, 1) - - encoder_output, encoder_output_lengths = self.encoder.get_audio_features(mel, mel_input_lengths) - encoder_max_input_length = torch.max(encoder_output_lengths).item() - output_ids = self.decoder.generate(decoder_input_ids, - encoder_output, - encoder_max_input_length, - encoder_output_lengths, - self.tokenizer.eot, - max_new_tokens=max_new_tokens, - num_beams=num_beams) + if self.use_py_session: + encoder_output, encoder_output_lengths = self.encoder.get_audio_features(mel, mel_input_lengths) + encoder_max_input_length = torch.max(encoder_output_lengths).item() + output_ids = self.decoder.generate(decoder_input_ids, + encoder_output, + encoder_max_input_length, + encoder_output_lengths, + self.tokenizer.eot, + max_new_tokens=max_new_tokens, + num_beams=num_beams) + else: + with torch.no_grad(): + if isinstance(mel, list): + mel = [ + m.transpose(1, 2).type( + str_dtype_to_torch("float16")).squeeze(0) + for m in mel + ] + else: + mel = mel.transpose(1, 2) + outputs = self.model_runner_cpp.generate( + batch_input_ids=decoder_input_ids, + encoder_input_features=mel, + encoder_output_lengths=mel_input_lengths // 2, + max_new_tokens=max_new_tokens, + end_id=self.tokenizer.eot, + pad_id=self.tokenizer.eot, + num_beams=num_beams, + output_sequence_lengths=True, + return_dict=True) + torch.cuda.synchronize() + output_ids = outputs['output_ids'].cpu().numpy().tolist() texts = [] for i in range(len(output_ids)): text = self.tokenizer.decode(output_ids[i][0]).strip() @@ -379,7 +425,8 @@ class WhisperTRTLLM(object): batch_size=1, num_beams=1, padding_strategy="max", - ): + max_new_tokens=96, + ): mel = mel.type(str_dtype_to_torch(dtype)) mel = mel.unsqueeze(0) # repeat the mel spectrogram to match the batch size @@ -393,7 +440,13 @@ class WhisperTRTLLM(object): dtype=torch.int32, device=mel.device) - predictions = self.process_batch(mel, features_input_lengths, text_prefix, num_beams) + predictions = self.process_batch( + mel, + features_input_lengths, + text_prefix, + num_beams, + max_new_tokens=max_new_tokens + ) prediction = predictions[0] # remove all special tokens in the prediction From 444a1df74086d6ad7a9da4675d6ada2878a76d6b Mon Sep 17 00:00:00 2001 From: Emmanuel Ferdman Date: Fri, 25 Apr 2025 00:30:08 -0700 Subject: [PATCH 2/2] Resolve daemon warnings for threading methods Signed-off-by: Emmanuel Ferdman --- whisper_live/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/whisper_live/client.py b/whisper_live/client.py index 4fc9281..dd41323 100644 --- a/whisper_live/client.py +++ b/whisper_live/client.py @@ -94,7 +94,7 @@ class Client: # start websocket client in a thread self.ws_thread = threading.Thread(target=self.client_socket.run_forever) - self.ws_thread.setDaemon(True) + self.ws_thread.daemon = True self.ws_thread.start() self.transcript = []