import time import math import os import re from copy import deepcopy import librosa import numpy as np from omegaconf import ListConfig, OmegaConf, open_dict import nemo.collections.asr as nemo_asr from nemo.collections.asr.models import SortformerEncLabelModel, ASRModel from nemo.collections.asr.parts.submodules.subsampling import FeatureStacking from nemo.collections.asr.parts.utils.streaming_utils import CacheAwareStreamingAudioBuffer from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis import torch from nemo.collections.asr.parts.utils.multispk_transcribe_utils import ( SpeakerTaggedASR, configure_diar_streaming, validate_feature_frame_strides, ) from typing import List, Optional, Union from dataclasses import dataclass, field ASR_LOCAL_MODEL_EXTENSIONS = (".nemo",) DIAR_LOCAL_MODEL_EXTENSIONS = (".nemo", ".ckpt") LANG_TAG_PATTERN = re.compile(r"^(?:<\|[^|]+?\|>\s*)+") def resolve_model_reference(model_reference: str, model_label: str, local_extensions) -> str: """Resolve a local checkpoint path or preserve a NeMo pretrained model identifier.""" if model_reference is None or not model_reference.strip(): raise ValueError(f"{model_label} model reference must not be empty") model_reference = model_reference.strip() expanded_reference = os.path.expanduser(model_reference) if os.path.isfile(expanded_reference): extension = os.path.splitext(expanded_reference)[1].lower() if extension not in local_extensions: raise ValueError( f"Unsupported local {model_label.lower()} model file: {model_reference}. " f"Expected one of: {', '.join(local_extensions)}" ) return expanded_reference if os.path.exists(expanded_reference): raise ValueError(f"{model_label} model reference must point to a checkpoint file: {model_reference}") extension = os.path.splitext(expanded_reference)[1].lower() looks_like_local_path = ( extension in local_extensions or model_reference.startswith(("/", "./", "../", "~")) ) if looks_like_local_path: raise FileNotFoundError(f"{model_label} model not found at: {model_reference}") # NeMo resolves identifiers containing "/" from Hugging Face Hub and bare # identifiers from its registered pretrained-model catalog. return model_reference def format_time(seconds): minutes = math.floor(seconds / 60) sec = seconds % 60 return f"{minutes}:{sec:05.2f}" def extract_transcriptions(hyps): """ The transcribed_texts returned by CTC and RNNT models are different. This method would extract and return the text section of the hypothesis. """ if isinstance(hyps[0], Hypothesis): transcriptions = [] for hyp in hyps: transcriptions.append(hyp.text) else: transcriptions = hyps return transcriptions def configure_asr_for_multitalker_streaming(config, asr_model): """Apply ASR-model-specific settings for the shared multitalker streaming path.""" if config.parallel_speaker_strategy and not config.masked_asr and not hasattr(asr_model, "set_speaker_targets"): raise ValueError( "parallel_speaker_strategy=True with masked_asr=False requires an ASR model that supports " "speaker-target injection. Use masked_asr=True for conventional streaming ASR models." ) if hasattr(asr_model, "set_inference_prompt"): target_lang = config.target_lang or "auto" asr_model.set_inference_prompt(target_lang) if config.strip_lang_tags and hasattr(asr_model, "cfg"): with open_dict(asr_model.cfg.decoding): asr_model.cfg.decoding.strip_lang_tags = True asr_model.cfg.decoding.lang_tag_pattern = config.lang_tag_pattern or LANG_TAG_PATTERN.pattern @dataclass class MultitalkerTranscriptionConfig: """ Configuration for Multi-talker transcription with an ASR model and a diarization model. """ # Required configs diar_model: Optional[str] = None # Local checkpoint or pretrained model identifier diar_pretrained_name: Optional[str] = None # Name of a pretrained model max_num_of_spks: Optional[int] = 4 # maximum number of speakers parallel_speaker_strategy: bool = True # whether to use parallel speaker strategy masked_asr: bool = False # whether to use masked ASR mask_preencode: bool = False # whether to mask preencode or mask features cache_gating: bool = True # whether to use cache gating cache_gating_buffer_size: int = 2 # buffer size for cache gating single_speaker_mode: bool = False # whether to use single speaker mode displayed_transcript_lines: int = 10 # number of transcript lines to display # General configs session_len_sec: float = -1 # End-to-end diarization session length in seconds num_workers: int = 8 random_seed: Optional[int] = None # seed number going to be used in seed_everything() log: bool = False # If True,log will be printed precision: str = "bf16" # Streaming diarization configs streaming_mode: bool = True # If True, streaming diarization will be used. spkcache_len: Optional[int] = None spkcache_update_period: int = 222 fifo_len: int = 264 diar_right_context: int = 0 # If `cuda` is a negative number, inference will be on CPU only. cuda: Optional[int] = None allow_mps: bool = False # allow to select MPS device (Apple Silicon M-series GPU) matmul_precision: str = "highest" # Literal["highest", "high", "medium"] # ASR Configs asr_model: Optional[str] = None # Local checkpoint or pretrained model identifier device: str = 'cuda' audio_file: Optional[str] = None manifest_file: Optional[str] = None use_amp: bool = True debug_mode: bool = False batch_size: int = 32 online_normalization: bool = False output_path: Optional[str] = None pad_and_drop_preencoded: bool = False set_decoder: Optional[str] = None # ["ctc", "rnnt"] att_context_size: Optional[List[int]] = field(default_factory=lambda: [56, 13]) target_lang: Optional[str] = "auto" strip_lang_tags: bool = True lang_tag_pattern: Optional[str] = LANG_TAG_PATTERN.pattern generate_realtime_scripts: bool = True # This should be True for huggingface space word_window: int = 50 sent_break_sec: float = 1.0 fix_prev_words_count: int = 5 update_prev_words_sentence: int = 5 left_frame_shift: int = -1 right_frame_shift: int = 0 min_sigmoid_val: float = 1e-2 discarded_frames: int = 8 print_time: bool = True print_sample_indices: List[int] = field(default_factory=lambda: [0]) colored_text: bool = True real_time_mode: bool = False print_path: Optional[str] = None ignored_initial_frame_steps: int = 5 verbose: bool = False feat_len_sec: float = 0.01 finetune_realtime_ratio: float = 0.01 spk_supervision: str = "diar" # ["diar", "rttm"] binary_diar_preds: bool = False # Microphone streaming configs sample_rate: int = 16000 microphone_streaming: bool = False microphone_chunk_size: int = 0 # Number of samples per chunk microphone_hop_length: int = 0 # Hop length for sliding window # Mix config deploy_mode: bool = True # Audio pre-loading configs audio_pre_loading: bool = True class StreamingSession: def __init__(self): self.config: Optional[MultitalkerTranscriptionConfig] = None self.asr_model: Optional[ASRModel] = None self.diar_model: Optional[SortformerEncLabelModel] = None self.streaming_buffer = None self.multispk_asr_streamer: Optional[SpeakerTaggedASR] = None self.transcribed_speaker_texts = "" self.step_num = 0 self.feat_frame_count = 0 self._max_num_of_spks = 4 self.chunk_audio_buffer = np.array([]) self.all_chunk_audio = np.array([]) self.session_start_time = None self.silence_level = 1e-4 self.silence_second = 0.5 self._is_result_buffer_called = False self._audio_pre_loading = True def load_asr_model( self, asr_model: str, device: Union[str, torch.device] = 'cuda', reload: bool = False, ): asr_model = resolve_model_reference( asr_model, model_label="ASR", local_extensions=ASR_LOCAL_MODEL_EXTENSIONS, ) if not reload: if os.path.isfile(asr_model): self.asr_model = nemo_asr.models.ASRModel.restore_from(restore_path=asr_model) else: self.asr_model = nemo_asr.models.ASRModel.from_pretrained( model_name=asr_model, map_location=device, ) self.asr_model.eval() self.asr_model = self.asr_model.to(device) # configure the decoding config decoding_cfg = self.asr_model.cfg.decoding with open_dict(decoding_cfg): decoding_cfg.strategy = "greedy_batch" decoding_cfg.preserve_alignments = False decoding_cfg.strip_lang_tags = self.config.strip_lang_tags decoding_cfg.lang_tag_pattern = self.config.lang_tag_pattern or LANG_TAG_PATTERN.pattern if hasattr(self.asr_model, 'joint'): # if an RNNT model decoding_cfg.greedy.max_symbols = 10 decoding_cfg.fused_batch_size = -1 self.asr_model.change_decoding_strategy(decoding_cfg) if self.config.att_context_size: self.asr_model.encoder.set_default_att_context_size(self.config.att_context_size) configure_asr_for_multitalker_streaming(self.config, self.asr_model) def load_diar_model( self, diar_model: str, device: Union[str, torch.device], reload: bool = False, ): diar_model = resolve_model_reference( diar_model, model_label="Diarization", local_extensions=DIAR_LOCAL_MODEL_EXTENSIONS, ) if not reload: if os.path.isfile(diar_model) and diar_model.endswith(".ckpt"): self.diar_model = SortformerEncLabelModel.load_from_checkpoint(checkpoint_path=diar_model, map_location=device, strict=False) elif os.path.isfile(diar_model): self.diar_model = SortformerEncLabelModel.restore_from(restore_path=diar_model, map_location=device) else: self.diar_model = SortformerEncLabelModel.from_pretrained( model_name=diar_model, map_location=device, ) self.diar_model.eval() self.diar_model = self.diar_model.to(device) if ( self.config.precision.startswith("bf16") and self.diar_model.device.type == "cuda" and torch.cuda.is_bf16_supported() ): self.diar_model = self.diar_model.to(dtype=torch.bfloat16) validate_feature_frame_strides(asr_model=self.asr_model, diar_model=self.diar_model) diar_chunk_len = ( self.asr_model.encoder.streaming_cfg.valid_out_len + self.asr_model.encoder.streaming_cfg.cache_drop_size ) configure_diar_streaming( diar_model=self.diar_model, cfg=self.config, output_subsampling_factor=self.asr_model.encoder.subsampling_factor, diar_chunk_len=diar_chunk_len, ) self.config.spkcache_len = int(self.diar_model.sortformer_modules.spkcache_len) def add_file_to_streaming_buffer(self, audio_file_path: str): # Load audio file at 16000 Hz using librosa (handles resampling automatically) audio, SR = librosa.load(audio_file_path, sr=self.config.sample_rate, mono=True) # Add very short silence (0.5 second) at the beginning of the audio for better streaming silence = np.random.randn(int(SR * self.silence_second)) * self.silence_level audio = np.concatenate([silence, audio]) # Preprocess and append to buffer self.streaming_buffer.append_audio(audio, stream_id=-1) # Create iterator for streaming self.streaming_buffer_iter = iter(self.streaming_buffer) def _setup_streaming_buffer(self): streaming_buffer = CacheAwareStreamingAudioBuffer( model=self.asr_model, online_normalization=False, pad_and_drop_preencoded=self.config.pad_and_drop_preencoded, ) return streaming_buffer def reset_buffer(self): # if self.streaming_buffer.buffer_idx != 0: print(f"[DEBUG] reset_buffer() Called at buffer_idx {self.streaming_buffer.buffer_idx}") self.streaming_buffer.reset_buffer() self.step_num = 0 self.feat_frame_count = 0 self._reset_microphone_audio_buffer() self.all_chunk_audio = np.array([]) self.session_start_time = None self._is_result_buffer_called = True self.transcribed_speaker_texts = "" def setup_streaming_session(self, config: MultitalkerTranscriptionConfig, reload: bool = False, reset_buffer: bool = False): """Setup streaming parameters based on model configuration""" self.transcribed_speaker_texts = "" self.config = OmegaConf.structured(config) self.load_asr_model(self.config.asr_model, self.config.device, reload=reload) self._configure_microphone_geometry() config.microphone_chunk_size = self.config.microphone_chunk_size config.microphone_hop_length = self.config.microphone_hop_length self.load_diar_model(self.config.diar_model, self.config.device, reload=reload) config.spkcache_len = self.config.spkcache_len self.config.max_num_of_spks = min( self.config.max_num_of_spks, int(self.diar_model._cfg.max_num_of_spks), ) config.max_num_of_spks = self.config.max_num_of_spks if isinstance(self.diar_model.encoder.pre_encode, FeatureStacking): self.config.pad_and_drop_preencoded = True # Always create new SpeakerTaggedASR instance to ensure clean state self.multispk_asr_streamer = SpeakerTaggedASR(self.config, self.asr_model, self.diar_model) print(f"[DEBUG] setup_streaming_session() Called, >>>>>> creating new SpeakerTaggedASR instance...") self.streaming_buffer = self._setup_streaming_buffer() self._max_num_of_spks = self.config.max_num_of_spks self._audio_pre_loading = self.config.audio_pre_loading self._displayed_transcript_lines = self.config.displayed_transcript_lines self._reset_microphone_audio_buffer() if reset_buffer: self.reset_buffer() def _configure_microphone_geometry(self): if not self.config.microphone_streaming: return streaming_cfg = self.asr_model.encoder.streaming_cfg feature_stride = float(self.asr_model.cfg.preprocessor.window_stride) hop_feature_frames = streaming_cfg.valid_out_len * self.asr_model.encoder.subsampling_factor hop_samples = round(hop_feature_frames * feature_stride * self.config.sample_rate) pre_encode_cache_size = streaming_cfg.pre_encode_cache_size if isinstance(pre_encode_cache_size, (list, tuple, ListConfig)): pre_encode_cache_size = pre_encode_cache_size[-1] cache_samples = round(pre_encode_cache_size * feature_stride * self.config.sample_rate) self.config.microphone_hop_length = int(hop_samples) self.config.microphone_chunk_size = int(hop_samples + cache_samples) def _reset_microphone_audio_buffer(self): if self.config is None or not self.config.microphone_streaming: self.chunk_audio_buffer = np.array([]) return cache_samples = self.config.microphone_chunk_size - self.config.microphone_hop_length self.chunk_audio_buffer = np.zeros(cache_samples, dtype=np.float32) def _get_left_right_offset(self, step_num): if step_num == 0: self.session_start_time = time.time() left_offset = 0 right_offset = 0 else: left_offset = 8 right_offset = 0 return left_offset, right_offset def _get_transcribed_texts(self): """ Extract transcribed texts from the multispk_asr_streamer's internal state. Returns a formatted string with speaker-tagged transcriptions. """ if self.multispk_asr_streamer is None: return "" # Get the word and timestamp sequences from the streamer word_and_ts_seq = self.multispk_asr_streamer._word_and_ts_seq transcriptions = [] for _uniq_id, data in word_and_ts_seq.items(): if data.get('sentences') is not None: for sentence in data['sentences']: speaker = sentence.get('speaker', 'Unknown') text = sentence.get('words', '').strip() if text: transcriptions.append(f"Speaker {speaker}: {text}") return "\n".join(transcriptions) if transcriptions else "" def _get_diar_visualization_state(self): """Collect retained cache predictions and fresh FIFO/chunk predictions for visualization.""" if ( self.multispk_asr_streamer is None or self.multispk_asr_streamer.instance_manager is None or self.multispk_asr_streamer.instance_manager.diar_states is None ): return None diar_state = self.multispk_asr_streamer.instance_manager.diar_states streaming_state = diar_state.streaming_state if streaming_state is None: return None def get_valid_length(embeddings, lengths): if lengths is not None: return int(lengths[0].item()) return 0 if embeddings is None else embeddings.shape[1] def trim_predictions(predictions, valid_length): if predictions is None: return None return predictions[:, :valid_length].detach() spkcache_preds = streaming_state.spkcache_preds fifo_preds = streaming_state.fifo_preds chunk_capacity = self.multispk_asr_streamer._nframes_per_chunk chunk_preds = ( None if diar_state.diar_pred_out_stream is None else diar_state.diar_pred_out_stream[:, -chunk_capacity:] ) retained_chunk_frames = 0 if fifo_preds is not None and chunk_preds is not None: retained_chunk_frames = min(fifo_preds.shape[1], chunk_preds.shape[1]) fifo_preds = fifo_preds[:, : fifo_preds.shape[1] - retained_chunk_frames] spkcache_valid = get_valid_length(streaming_state.spkcache, streaming_state.spkcache_lengths) fifo_valid = 0 if fifo_preds is None else fifo_preds.shape[1] chunk_valid = 0 if chunk_preds is None else chunk_preds.shape[1] return { "spkcache": trim_predictions(spkcache_preds, spkcache_valid), "fifo": trim_predictions(fifo_preds, fifo_valid), "chunk": None if chunk_preds is None else chunk_preds.detach(), "valid_frames": { "spkcache": spkcache_valid, "fifo": fifo_valid, "chunk": chunk_valid, }, "capacities": { "spkcache": self.config.spkcache_len, "fifo": self.config.fifo_len, "chunk": chunk_capacity, }, } @torch.inference_mode() def process_audio_chunk(self, chunk_audio, chunk_lengths, is_buffer_empty=None): """ Process an audio chunk using the SpeakerTaggedASR streamer. The streamer handles all internal state management through its instance_manager. Args: chunk_audio: torch.Tensor of shape (batch, channels, time) chunk_lengths: torch.Tensor of shape (batch,) with lengths """ # Determine drop_extra_pre_encoded based on offsets # drop_extra_pre_encoded = left_offset drop_extra_pre_encoded = ( 0 if self.step_num == 0 and not self.config.pad_and_drop_preencoded else self.asr_model.encoder.streaming_cfg.drop_extra_pre_encoded ) if is_buffer_empty is None: is_buffer_empty = self.streaming_buffer.buffer is not None and self.streaming_buffer.is_buffer_empty() use_bf16 = ( self.config.precision.startswith("bf16") and self.asr_model.device.type == "cuda" and torch.cuda.is_bf16_supported() ) # Call the appropriate streaming method based on config with torch.amp.autocast( device_type=self.asr_model.device.type, dtype=torch.bfloat16 if use_bf16 else self.asr_model.dtype, enabled=self.config.use_amp and use_bf16, ): if self.config.parallel_speaker_strategy: # Parallel streaming: multiple ASR instances for multiple speakers transcribed_speaker_texts = self.multispk_asr_streamer.perform_parallel_streaming_stt_spk( step_num=self.step_num, chunk_audio=chunk_audio, chunk_lengths=chunk_lengths, is_buffer_empty=is_buffer_empty, drop_extra_pre_encoded=drop_extra_pre_encoded, ) else: # Serial streaming: single ASR instance for all speakers transcribed_speaker_texts = self.multispk_asr_streamer.perform_serial_streaming_stt_spk( step_num=self.step_num, chunk_audio=chunk_audio, chunk_lengths=chunk_lengths, is_buffer_empty=is_buffer_empty, drop_extra_pre_encoded=drop_extra_pre_encoded, ) # Extract transcriptions from the instance manager's state if transcribed_speaker_texts is not None: self.transcribed_speaker_texts = transcribed_speaker_texts[0] self.feat_frame_count += (chunk_audio.shape[-1] - self.config.discarded_frames) if self.config.log: print(f'[Deploy Mode: {self.config.deploy_mode}] transcribed_speaker_texts:', self.transcribed_speaker_texts) self.step_num += 1 return self._get_diar_visualization_state() def flush_microphone_buffer(self): """Process microphone audio that has not yet filled a complete streaming hop.""" if not self.config.microphone_streaming: return None cache_samples = self.config.microphone_chunk_size - self.config.microphone_hop_length if len(self.chunk_audio_buffer) <= cache_samples: return None frame = self.chunk_audio_buffer.copy() self.chunk_audio_buffer = ( frame[-cache_samples:].copy() if cache_samples > 0 else np.array([], dtype=np.float32) ) chunk_audio, chunk_lengths = self.streaming_buffer.preprocess_audio(frame) chunk_audio = chunk_audio[:, :, :chunk_lengths[0]] return self.process_audio_chunk(chunk_audio, chunk_lengths, is_buffer_empty=True) def process_microphone_chunk(self, chunk_audio, sr): # Convert to float32 and normalize if chunk_audio.dtype == np.int16: chunk_audio = chunk_audio.astype(np.float32) / 32768.0 elif chunk_audio.dtype == np.int32: chunk_audio = chunk_audio.astype(np.float32) / 2147483648.0 if sr != self.config.sample_rate: chunk_audio = librosa.resample(chunk_audio, orig_sr=sr, target_sr=self.config.sample_rate) self.all_chunk_audio = np.concatenate([self.all_chunk_audio, chunk_audio]) self.chunk_audio_buffer = np.concatenate([self.chunk_audio_buffer, chunk_audio]) processed_flag = False diar_pred_out = None while len(self.chunk_audio_buffer) >= self.config.microphone_chunk_size: frame = self.chunk_audio_buffer[:self.config.microphone_chunk_size] print(f"[DEBUG] process_microphone_chunk() L0: frame.shape: {frame.shape}") self.chunk_audio_buffer = self.chunk_audio_buffer[self.config.microphone_hop_length:] chunk_audio, chunk_lengths = self.streaming_buffer.preprocess_audio(frame) print(f"[DEBUG] process_microphone_chunk() L1: original chunk_audio.shape: {chunk_audio.shape}, chunk_lengths: {chunk_lengths} frame.shape: {frame.shape}") chunk_audio = chunk_audio[:, :, :chunk_lengths[0]] print(f"[DEBUG] process_microphone_chunk() L2: self.config.microphone_hop_length {self.config.microphone_hop_length} self.config.microphone_chunk_size {self.config.microphone_chunk_size}") print(f"[DEBUG] process_microphone_chunk() L3: Truncated chunk_audio.shape: {chunk_audio.shape}") print(f"[DEBUG] self._max_num_of_spks: {self._max_num_of_spks} and self.multispk_asr_streamer._max_num_of_spks: {self.multispk_asr_streamer._max_num_of_spks}") if self.step_num == 0: self.session_start_time = time.time() diar_pred_out = self.process_audio_chunk(chunk_audio, chunk_lengths, is_buffer_empty=False) processed_flag = True if processed_flag: print(f">>>> Finished processing microphone chunk: {chunk_audio.shape} {sr}") else: print(f">>>> Skipping processing microphone chunk: {chunk_audio.shape} {sr}") return diar_pred_out