""" Voxtral Mini Realtime streaming backend using HuggingFace Transformers. Uses VoxtralRealtimeForConditionalGeneration with a background generate thread and queue-based audio feeding for real-time streaming transcription. Supports CUDA, CPU, and MPS devices. """ import logging import queue import sys import threading import time from typing import List, Optional, Tuple import numpy as np from whisperlivekit.timed_objects import ASRToken, Transcript logger = logging.getLogger(__name__) class VoxtralHFStreamingASR: """Voxtral model holder using HuggingFace Transformers.""" sep = " " def __init__(self, logfile=sys.stderr, **kwargs): import torch from transformers import ( AutoProcessor, VoxtralRealtimeForConditionalGeneration, ) self.logfile = logfile self.transcribe_kargs = {} lan = kwargs.get("lan", "auto") self.original_language = None if lan == "auto" else lan DEFAULT_MODEL = "mistralai/Voxtral-Mini-4B-Realtime-2602" model_path = kwargs.get("model_dir") or kwargs.get("model_path") if not model_path: model_size = kwargs.get("model_size", "") if model_size and ("/" in model_size or model_size.startswith(".")): model_path = model_size else: model_path = DEFAULT_MODEL t = time.time() logger.info(f"Loading Voxtral model '{model_path}' via HF Transformers...") self.processor = AutoProcessor.from_pretrained(model_path) self.model = VoxtralRealtimeForConditionalGeneration.from_pretrained( model_path, torch_dtype=torch.bfloat16, device_map="auto", ) logger.info(f"Voxtral HF model loaded in {time.time() - t:.2f}s on {self.model.device}") self.backend_choice = "voxtral" self.tokenizer = None # sentence tokenizer — not needed for streaming def transcribe(self, audio): pass class VoxtralHFStreamingOnlineProcessor: """ Online processor for Voxtral streaming ASR via HuggingFace Transformers. Uses a background thread running model.generate() with a queue-based input_features_generator and TextIteratorStreamer for real-time output. Each decoded token corresponds to ~80ms of audio. """ SAMPLING_RATE = 16000 def __init__(self, asr: VoxtralHFStreamingASR, logfile=sys.stderr): self.asr = asr self.logfile = logfile self.end = 0.0 self.buffer = [] self.audio_buffer = np.array([], dtype=np.float32) processor = asr.processor self._first_chunk_samples = processor.num_samples_first_audio_chunk self._chunk_samples = processor.num_samples_per_audio_chunk self._chunk_step = processor.raw_audio_length_per_tok # num_right_pad_tokens is a method in some transformers versions, a property in others n_right_pad = processor.num_right_pad_tokens if callable(n_right_pad): n_right_pad = n_right_pad() self._right_pad_samples = int(n_right_pad * processor.raw_audio_length_per_tok) self._seconds_per_token = processor.raw_audio_length_per_tok / self.SAMPLING_RATE self._reset_state() logger.info( f"[voxtral-hf] Initialized. first_chunk={self._first_chunk_samples} samples, " f"chunk={self._chunk_samples}, step={self._chunk_step}, " f"right_pad={self._right_pad_samples}" ) def _reset_state(self): self._pending_audio = np.zeros(0, dtype=np.float32) self._audio_queue: queue.Queue = queue.Queue() self._streamer_texts: List[str] = [] self._generate_thread: Optional[threading.Thread] = None self._generate_started = False self._generate_finished = False self._generate_error: Optional[Exception] = None # Text accumulation and word extraction self._accumulated_text = "" self._n_text_tokens_received = 0 self._n_audio_tokens_fed = 0 self._n_committed_words = 0 self._global_time_offset = 0.0 # Lock for text state accessed from both generate thread and main thread self._text_lock = threading.Lock() # ── Interface methods ── def insert_audio_chunk(self, audio: np.ndarray, audio_stream_end_time: float): self.end = audio_stream_end_time self._pending_audio = np.append(self._pending_audio, audio) self.audio_buffer = self._pending_audio def process_iter(self, is_last=False) -> Tuple[List[ASRToken], float]: try: return self._process_iter_inner(is_last) except Exception as e: logger.warning(f"[voxtral-hf] process_iter exception: {e}", exc_info=True) return [], self.end def get_buffer(self) -> Transcript: """Return all uncommitted text as buffer. Drains the streamer first so late-arriving tokens (common on slower devices like MPS) are picked up even between audio chunks. """ self._drain_streamer() with self._text_lock: text = self._accumulated_text if not text: return Transcript(start=None, end=None, text="") words = text.split() uncommitted = words[self._n_committed_words:] if uncommitted: return Transcript(start=self.end, end=self.end, text=" ".join(uncommitted)) return Transcript(start=None, end=None, text="") def start_silence(self) -> Tuple[List[ASRToken], float]: """Flush all uncommitted words when silence starts. Feeds right-padding (silence) so the model has enough future context to emit the last few tokens, then drains repeatedly until the model has finished producing text. Without right-padding the model holds back the last few words because it hasn't seen enough audio yet. """ if not self._generate_started or self._generate_finished: self._drain_streamer() words = self._flush_all_pending_words() logger.info(f"[voxtral-hf] start_silence (no thread): flushed {len(words)} words") return words, self.end # Feed any remaining real audio self._feed_pending_audio() # Add right-padding so the model can decode trailing tokens. # Don't count these toward _n_audio_tokens_fed — they're not # real audio and shouldn't affect word timestamp calculations. if self._right_pad_samples > 0: right_pad = np.zeros(self._right_pad_samples, dtype=np.float32) self._pending_audio = np.append(self._pending_audio, right_pad) saved_count = self._n_audio_tokens_fed self._feed_pending_audio() self._n_audio_tokens_fed = saved_count # Drain in a loop: the model may still be processing right-padding # chunks after the first drain returns. Keep draining until no new # text appears for two consecutive rounds. all_words: List[ASRToken] = [] for _ in range(5): # at most 5 drain+flush cycles self._drain_streamer_blocking(timeout=5.0) batch = self._flush_all_pending_words() all_words.extend(batch) if not batch: break # no new text — model has caught up logger.info(f"[voxtral-hf] start_silence: flushed {len(all_words)} words") return all_words, self.end def end_silence(self, silence_duration: float, offset: float): self._global_time_offset += silence_duration self.end += silence_duration def new_speaker(self, change_speaker): self.start_silence() def warmup(self, audio, init_prompt=""): pass def finish(self) -> Tuple[List[ASRToken], float]: """Flush remaining audio with right-padding and stop the generate thread.""" # Add right-padding so the model can finish decoding if self._right_pad_samples > 0: right_pad = np.zeros(self._right_pad_samples, dtype=np.float32) self._pending_audio = np.append(self._pending_audio, right_pad) # Feed remaining audio if self._generate_started and not self._generate_finished: self._feed_pending_audio() # Signal end of audio self._audio_queue.put(None) # Wait for generate to finish if self._generate_thread is not None: self._generate_thread.join(timeout=30.0) elif not self._generate_started and len(self._pending_audio) >= self._first_chunk_samples: # Never started but have enough audio — start and immediately finish self._start_generate_thread() self._feed_pending_audio() self._audio_queue.put(None) if self._generate_thread is not None: self._generate_thread.join(timeout=30.0) self._drain_streamer() words = self._flush_all_pending_words() logger.info(f"[voxtral-hf] finish: flushed {len(words)} words") return words, self.end # ── Generate thread management ── def _start_generate_thread(self): """Start model.generate() in a background thread with streaming.""" import torch from transformers import TextIteratorStreamer processor = self.asr.processor model = self.asr.model # Extract first chunk first_chunk_audio = self._pending_audio[:self._first_chunk_samples] self._pending_audio = self._pending_audio[self._first_chunk_samples:] # First chunk covers multiple audio tokens self._n_audio_tokens_fed += max(1, self._first_chunk_samples // self._chunk_step) first_inputs = processor( first_chunk_audio, is_streaming=True, is_first_audio_chunk=True, return_tensors="pt", ) first_inputs = first_inputs.to(model.device, dtype=model.dtype) streamer = TextIteratorStreamer( processor.tokenizer, skip_prompt=True, skip_special_tokens=True, ) self._streamer = streamer audio_queue = self._audio_queue def input_features_gen(): yield first_inputs.input_features while True: chunk_audio = audio_queue.get() if chunk_audio is None: break inputs = processor( chunk_audio, is_streaming=True, is_first_audio_chunk=False, return_tensors="pt", ) inputs = inputs.to(model.device, dtype=model.dtype) yield inputs.input_features def run_generate(): try: with torch.no_grad(): # Pass generator as input_features — the model detects GeneratorType # and internally converts it to input_features_generator generate_kwargs = { k: v for k, v in first_inputs.items() if k != "input_features" } model.generate( input_features=input_features_gen(), streamer=streamer, **generate_kwargs, ) except Exception as e: logger.error(f"[voxtral-hf] generate error: {e}", exc_info=True) self._generate_error = e finally: self._generate_finished = True self._generate_thread = threading.Thread(target=run_generate, daemon=True) self._generate_thread.start() self._generate_started = True logger.info("[voxtral-hf] generate thread started") def _feed_pending_audio(self): """Convert pending audio into properly-sized chunks for the generator.""" chunk_size = self._chunk_samples step_size = self._chunk_step while len(self._pending_audio) >= chunk_size: chunk = self._pending_audio[:chunk_size] self._audio_queue.put(chunk) self._pending_audio = self._pending_audio[step_size:] self._n_audio_tokens_fed += 1 self.audio_buffer = self._pending_audio def _drain_streamer(self): """Non-blocking drain of all available text from the streamer.""" if not self._generate_started: return text_queue = self._streamer.text_queue while True: try: text_fragment = text_queue.get_nowait() except queue.Empty: break if text_fragment is None: self._generate_finished = True break if text_fragment: with self._text_lock: self._accumulated_text += text_fragment self._n_text_tokens_received += 1 def _drain_streamer_blocking(self, timeout=30.0): """Blocking drain: wait for the generate thread to process all queued audio and produce the corresponding text. Polls the text queue while the audio queue has items (model still processing). Once the audio queue is empty, waits for trailing tokens, then returns. This is critical for start_silence(): without it, the non-blocking drain races with the generate thread and the last words get stuck. """ if not self._generate_started or self._generate_finished: self._drain_streamer() return text_queue = self._streamer.text_queue deadline = time.time() + timeout while time.time() < deadline: # Short poll while model is still processing queued audio; # longer wait once the audio queue is empty (trailing tokens). wait = 2.0 if self._audio_queue.empty() else 0.1 try: text_fragment = text_queue.get(timeout=wait) except queue.Empty: if self._audio_queue.empty(): break # Audio done + no text for 2s → fully caught up continue # Audio still queued, model still working if text_fragment is None: self._generate_finished = True break if text_fragment: with self._text_lock: self._accumulated_text += text_fragment self._n_text_tokens_received += 1 # ── Word extraction ── def _pos_to_time(self, token_position: int) -> float: """Convert token position to seconds.""" return token_position * self._seconds_per_token + self._global_time_offset def _extract_new_words(self) -> List[ASRToken]: """Extract complete words (all but the last, which may still be growing).""" with self._text_lock: text = self._accumulated_text if not text: return [] words = text.split() new_words: List[ASRToken] = [] n_words_total = len(words) n_audio_toks = max(self._n_audio_tokens_fed, 1) while len(words) > self._n_committed_words + 1: word = words[self._n_committed_words] word_idx = self._n_committed_words tok_start = int(word_idx / n_words_total * n_audio_toks) if n_words_total > 0 else 0 tok_end = int((word_idx + 1) / n_words_total * n_audio_toks) if n_words_total > 0 else 0 start_time = self._pos_to_time(tok_start) end_time = self._pos_to_time(tok_end) text_out = word if self._n_committed_words == 0 else " " + word new_words.append(ASRToken(start=start_time, end=end_time, text=text_out)) self._n_committed_words += 1 return new_words def _flush_all_pending_words(self) -> List[ASRToken]: """Flush ALL words including the last partial one.""" with self._text_lock: text = self._accumulated_text if not text: return [] words = text.split() new_words: List[ASRToken] = [] n_words_total = max(len(words), 1) n_audio_toks = max(self._n_audio_tokens_fed, 1) while self._n_committed_words < len(words): word = words[self._n_committed_words] word_idx = self._n_committed_words tok_start = int(word_idx / n_words_total * n_audio_toks) tok_end = int((word_idx + 1) / n_words_total * n_audio_toks) start_time = self._pos_to_time(tok_start) end_time = self._pos_to_time(tok_end) text_out = word if self._n_committed_words == 0 else " " + word new_words.append(ASRToken(start=start_time, end=end_time, text=text_out)) self._n_committed_words += 1 return new_words # ── Core processing ── def _process_iter_inner(self, is_last: bool) -> Tuple[List[ASRToken], float]: # Start generate thread when enough audio is buffered if not self._generate_started: if len(self._pending_audio) >= self._first_chunk_samples: self._start_generate_thread() self._feed_pending_audio() else: return [], self.end # Feed any new pending audio if self._generate_started and not self._generate_finished: self._feed_pending_audio() # If generate finished unexpectedly (EOS) but new audio arrived, restart if self._generate_finished and len(self._pending_audio) >= self._first_chunk_samples: self._drain_streamer() flush_words = self._flush_all_pending_words() # Reset for new utterance old_offset = self._global_time_offset self._reset_state() self._global_time_offset = old_offset self._start_generate_thread() self._feed_pending_audio() return flush_words, self.end # Drain available text from streamer self._drain_streamer() # Extract complete words new_words = self._extract_new_words() if new_words: logger.info(f"[voxtral-hf] returning {len(new_words)} words: {[w.text for w in new_words]}") self.buffer = [] return new_words, self.end