From 728e1f1290f6f4927d5c46fb15b78499df575bdd Mon Sep 17 00:00:00 2001 From: Quentin Fuxa Date: Tue, 12 Aug 2025 18:35:04 +0200 Subject: [PATCH] simulstreaming warmup is done for each instance of online, not for the backend --- whisperlivekit/core.py | 5 +- whisperlivekit/simul_whisper/backend.py | 26 ++++---- whisperlivekit/warmup.py | 62 ++++++++++++++++++ .../whisper_online.py | 65 +------------------ 4 files changed, 80 insertions(+), 78 deletions(-) create mode 100644 whisperlivekit/warmup.py diff --git a/whisperlivekit/core.py b/whisperlivekit/core.py index adfd96b..59443e5 100644 --- a/whisperlivekit/core.py +++ b/whisperlivekit/core.py @@ -1,10 +1,10 @@ try: - from whisperlivekit.whisper_streaming_custom.whisper_online import backend_factory, warmup_asr + from whisperlivekit.whisper_streaming_custom.whisper_online import backend_factory from whisperlivekit.whisper_streaming_custom.online_asr import VACOnlineASRProcessor, OnlineASRProcessor except ImportError: from .whisper_streaming_custom.whisper_online import backend_factory, warmup_asr from .whisper_streaming_custom.online_asr import VACOnlineASRProcessor, OnlineASRProcessor - +from whisperlivekit.warmup import warmup_asr, warmup_online from argparse import Namespace import sys @@ -130,6 +130,7 @@ def online_factory(args, asr, tokenizer, logfile=sys.stderr): asr, logfile=logfile, ) + warmup_online(online, args.warmup_file) elif args.vac: online = VACOnlineASRProcessor( args.min_chunk_size, diff --git a/whisperlivekit/simul_whisper/backend.py b/whisperlivekit/simul_whisper/backend.py index 4af7588..b8456e4 100644 --- a/whisperlivekit/simul_whisper/backend.py +++ b/whisperlivekit/simul_whisper/backend.py @@ -26,6 +26,7 @@ class SimulStreamingOnlineProcessor: self, asr, logfile=sys.stderr, + warmup_file=None ): self.asr = asr self.logfile = logfile @@ -121,6 +122,19 @@ class SimulStreamingOnlineProcessor: logger.exception(f"SimulStreaming processing error: {e}") return [], self.end + def warmup(self, audio, init_prompt=""): + """Warmup the SimulStreaming model.""" + try: + if isinstance(audio, np.ndarray): + audio = torch.from_numpy(audio).float() + self.model.insert_audio(audio) + self.model.infer(True) + self.model.refresh_segment(complete=True) + logger.info("SimulStreaming model warmed up successfully") + except Exception as e: + logger.exception(f"SimulStreaming warmup failed: {e}") + + class SimulStreamingASR(): """SimulStreaming backend with AlignAtt policy.""" sep = "" @@ -203,15 +217,3 @@ class SimulStreamingASR(): num_languages=self.model.model.num_languages, task="translate" ) - - # def warmup(self, audio, init_prompt=""): - # """Warmup the SimulStreaming model.""" - # try: - # if isinstance(audio, np.ndarray): - # audio = torch.from_numpy(audio).float() - # self.model.insert_audio(audio) - # self.model.infer(True) - # self.model.refresh_segment(complete=True) - # logger.info("SimulStreaming model warmed up successfully") - # except Exception as e: - # logger.exception(f"SimulStreaming warmup failed: {e}") diff --git a/whisperlivekit/warmup.py b/whisperlivekit/warmup.py new file mode 100644 index 0000000..008cace --- /dev/null +++ b/whisperlivekit/warmup.py @@ -0,0 +1,62 @@ + +import logging + +logger = logging.getLogger(__name__) + +def load_file(warmup_file=None, timeout=5): + import os + import tempfile + import librosa + + if warmup_file is None: + # Download JFK sample if not already present + jfk_url = "https://github.com/ggerganov/whisper.cpp/raw/master/samples/jfk.wav" + temp_dir = tempfile.gettempdir() + warmup_file = os.path.join(temp_dir, "whisper_warmup_jfk.wav") + + if not os.path.exists(warmup_file): + logger.debug(f"Downloading warmup file from {jfk_url}") + print(f"Downloading warmup file from {jfk_url}") + import time + import urllib.request + import urllib.error + import socket + + original_timeout = socket.getdefaulttimeout() + socket.setdefaulttimeout(timeout) + + start_time = time.time() + try: + urllib.request.urlretrieve(jfk_url, warmup_file) + logger.debug(f"Download successful in {time.time() - start_time:.2f}s") + except (urllib.error.URLError, socket.timeout) as e: + logger.warning(f"Download failed: {e}. Proceeding without warmup.") + return False + finally: + socket.setdefaulttimeout(original_timeout) + elif not warmup_file: + return False + + if not warmup_file or not os.path.exists(warmup_file) or os.path.getsize(warmup_file) == 0: + logger.warning(f"Warmup file {warmup_file} invalid or missing.") + return False + + try: + audio, sr = librosa.load(warmup_file, sr=16000) + except Exception as e: + logger.warning(f"Failed to load audio file: {e}") + return False + return audio + +def warmup_asr(asr, warmup_file=None, timeout=5): + """ + Warmup the ASR model by transcribing a short audio file. + """ + audio = load_file(warmup_file=None, timeout=5) + asr.warmup(audio) + logger.info("ASR model is warmed up") + +def warmup_online(online, warmup_file=None, timeout=5): + audio = load_file(warmup_file=None, timeout=5) + online.warmup(audio) + logger.warning("ASR is warmed up") \ No newline at end of file diff --git a/whisperlivekit/whisper_streaming_custom/whisper_online.py b/whisperlivekit/whisper_streaming_custom/whisper_online.py index 25349ac..db59027 100644 --- a/whisperlivekit/whisper_streaming_custom/whisper_online.py +++ b/whisperlivekit/whisper_streaming_custom/whisper_online.py @@ -107,67 +107,4 @@ def backend_factory(args): tokenizer = create_tokenizer(tgt_language) else: tokenizer = None - return asr, tokenizer - -def warmup_asr(asr, warmup_file=None, timeout=5): - """ - Warmup the ASR model by transcribing a short audio file. - """ - import os - import tempfile - - is_simulstreaming = hasattr(asr, 'warmup') and callable(getattr(asr, 'warmup')) - - if warmup_file is None: - # Download JFK sample if not already present - jfk_url = "https://github.com/ggerganov/whisper.cpp/raw/master/samples/jfk.wav" - temp_dir = tempfile.gettempdir() - warmup_file = os.path.join(temp_dir, "whisper_warmup_jfk.wav") - - if not os.path.exists(warmup_file): - logger.debug(f"Downloading warmup file from {jfk_url}") - print(f"Downloading warmup file from {jfk_url}") - import time - import urllib.request - import urllib.error - import socket - - original_timeout = socket.getdefaulttimeout() - socket.setdefaulttimeout(timeout) - - start_time = time.time() - try: - urllib.request.urlretrieve(jfk_url, warmup_file) - logger.debug(f"Download successful in {time.time() - start_time:.2f}s") - except (urllib.error.URLError, socket.timeout) as e: - logger.warning(f"Download failed: {e}. Proceeding without warmup.") - return False - finally: - socket.setdefaulttimeout(original_timeout) - elif not warmup_file: - return False - - if not warmup_file or not os.path.exists(warmup_file) or os.path.getsize(warmup_file) == 0: - logger.warning(f"Warmup file {warmup_file} invalid or missing.") - return False - - print(f"Warming up {'SimulStreaming' if is_simulstreaming else 'Whisper'} with {warmup_file}") - try: - import librosa - audio, sr = librosa.load(warmup_file, sr=16000) - except Exception as e: - logger.warning(f"Failed to load audio file: {e}") - return False - - try: - if is_simulstreaming: - asr.warmup(audio) - else: - asr.transcribe(audio) - - logger.info(f"{'SimulStreaming' if is_simulstreaming else 'Whisper'} is warmed up") - return True - - except Exception as e: - logger.warning(f"Warmup failed: {e}") - return False + return asr, tokenizer \ No newline at end of file