diff --git a/examples/speaker_tasks/diarization/README.md b/examples/speaker_tasks/diarization/README.md index 886abb73b8bd..d3ec6d248e03 100644 --- a/examples/speaker_tasks/diarization/README.md +++ b/examples/speaker_tasks/diarization/README.md @@ -86,6 +86,33 @@ diar_model.sortformer_modules.spkcache_update_period = 300 predicted_segments = diar_model.diarize(audio="/path/to/audio.wav", batch_size=1) ``` +For live mono waveforms, create a session with a fixed number of streams and push a padded float tensor at the model +sample rate together with its valid sample lengths. Each row owns independent preprocessing buffers and asynchronous +Sortformer cache state. A step can return no frames for a row until its configured chunk and right context are +available. Use a per-row ``is_final`` mask to flush streams independently; finalized rows remain in the batch with a +zero input length while other rows continue. Call ``reset()`` before reusing the session. The session disables dither +and feature padding while extracting each chunk, then applies the checkpoint's feature normalization over each +complete model input window. + +Use ``max_speakers`` to enable fewer than the checkpoint's maximum number of speaker channels. A scalar applies to +every stream, while a sequence or integer tensor supplies one limit per row. Disabled channels are zeroed in returned +probabilities and excluded from speaker-cache and FIFO state. Their raw activity is used only to prevent speech frames +from being added to the running silence profile. + +```python +import torch + +session = diar_model.create_streaming_session(batch_size=2, max_speakers=[2, 4]) +for audio_batch, audio_lengths, final_mask in audio_stream: + probabilities, probability_lengths = session.diarize_step( + audio_batch, + audio_chunk_lengths=audio_lengths, + is_final=final_mask, + ) + stream_0_probabilities = probabilities[0, : probability_lengths[0]] +session.reset() +``` + Diarization Error Rate (DER) with post-processing — all evaluations include overlapping speech: | Dataset | Collar | 30.4s latency | 10.0s latency | 1.04s latency | 0.32s latency | diff --git a/nemo/collections/asr/models/sortformer_diar_models.py b/nemo/collections/asr/models/sortformer_diar_models.py index f868a6e15581..58e299118b82 100644 --- a/nemo/collections/asr/models/sortformer_diar_models.py +++ b/nemo/collections/asr/models/sortformer_diar_models.py @@ -18,7 +18,7 @@ import os import random from collections import OrderedDict -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union import numpy as np import torch @@ -615,6 +615,28 @@ def process_signal(self, audio_signal, audio_signal_length): torch.cuda.empty_cache() return processed_signal, processed_signal_length + def create_streaming_session( + self, + batch_size: int = 1, + max_speakers: Optional[Union[int, Sequence[int], torch.Tensor]] = None, + ): + """Create an independent high-level raw-audio streaming session for a fixed batch of streams. + + The returned session accepts arbitrarily sized mono waveform chunks through ``diarize_step()`` and owns the + per-stream preprocessing buffers and batched asynchronous Sortformer cache state. + + Args: + batch_size: Fixed number of independent audio streams owned by the session. + max_speakers: Number of enabled speaker channels for every stream. A scalar applies to the complete batch; + a sequence or tensor supplies one value per row. By default, every model speaker channel is enabled. + + Returns: + SortformerStreamingSession: A new session bound to this model. + """ + from nemo.collections.asr.parts.utils.sortformer_utils import SortformerStreamingSession + + return SortformerStreamingSession(self, batch_size=batch_size, max_speakers=max_speakers) + def forward( self, audio_signal, @@ -938,6 +960,7 @@ def forward_streaming_step( drop_extra_pre_encoded=0, left_offset=0, right_offset=0, + async_streaming=None, ): """ One-step forward pass for diarization inference in streaming mode. @@ -966,6 +989,7 @@ def forward_streaming_step( drop_extra_pre_encoded (int): Number of leading pre-encoded frames to discard before streaming updates. left_offset (int): left offset for the current chunk right_offset (int): right offset for the current chunk + async_streaming (Optional[bool]): Override the model-level asynchronous streaming setting for this call. Returns: streaming_state (SortformerStreamingState): @@ -975,6 +999,11 @@ def forward_streaming_step( Tensor containing the updated total predicted speaker activity probabilities. Shape: (batch_size, cumulative pred length, num_speakers) """ + if async_streaming is None: + async_streaming = self.async_streaming + elif not isinstance(async_streaming, bool): + raise TypeError(f"async_streaming must be a boolean or None, got {type(async_streaming).__name__}") + chunk_pre_encode_embs, chunk_pre_encode_lengths = self._call_pre_encode( processed_signal, processed_signal_length ) @@ -983,7 +1012,7 @@ def forward_streaming_step( chunk_pre_encode_embs = chunk_pre_encode_embs[:, drop_extra_pre_encoded:, :] chunk_pre_encode_lengths = chunk_pre_encode_lengths - drop_extra_pre_encoded - if self.async_streaming: + if async_streaming: output_length = None if self.async_pad_to_max: output_length = ( @@ -1024,7 +1053,7 @@ def forward_streaming_step( spkcache_fifo_chunk_preds = self.sortformer_modules.downsample_preds( high_resolution_preds, self.upsample_factor ).detach() - if not self.async_streaming and streaming_state.spk_perm is not None: + if not async_streaming and streaming_state.spk_perm is not None: inv_spk_perm = torch.stack( [ torch.argsort(streaming_state.spk_perm[batch_index]) @@ -1037,11 +1066,22 @@ def forward_streaming_step( for batch_index in range(high_resolution_preds.shape[0]) ] ) + high_resolution_preds = self.sortformer_modules.apply_max_speakers_mask( + high_resolution_preds, streaming_state.max_speakers + ) spkcache_fifo_chunk_preds = self.sortformer_modules.apply_mask_to_preds( spkcache_fifo_chunk_preds, spkcache_fifo_chunk_fc_encoder_lengths ) - if self.async_streaming: + silence_profile_preds = None + if streaming_state.max_speakers is not None: + # Disabled channels must not enter output/cache state, but their activity still prevents a speech frame + # from being folded into the running silence embedding. + silence_profile_preds = spkcache_fifo_chunk_preds + spkcache_fifo_chunk_preds = self.sortformer_modules.apply_max_speakers_mask( + spkcache_fifo_chunk_preds, streaming_state.max_speakers + ) + if async_streaming: saved_spkcache_lengths = streaming_state.spkcache_lengths.clone() saved_fifo_lengths = streaming_state.fifo_lengths.clone() streaming_state, chunk_preds = self.sortformer_modules.streaming_update_async( @@ -1051,6 +1091,7 @@ def forward_streaming_step( preds=spkcache_fifo_chunk_preds, lc=lc_enc, rc=rc_enc, + silence_profile_preds=silence_profile_preds, ) if self.high_resolution: max_chunk_len = chunk_pre_encode_embs.shape[1] - lc_enc - rc_enc @@ -1072,6 +1113,7 @@ def forward_streaming_step( preds=spkcache_fifo_chunk_preds, lc=lc_enc, rc=rc_enc, + silence_profile_preds=silence_profile_preds, ) if self.high_resolution: chunk_len = chunk_pre_encode_embs.shape[1] - lc_enc - rc_enc diff --git a/nemo/collections/asr/modules/sortformer_modules.py b/nemo/collections/asr/modules/sortformer_modules.py index 0a19c3891cd2..2314f555d479 100644 --- a/nemo/collections/asr/modules/sortformer_modules.py +++ b/nemo/collections/asr/modules/sortformer_modules.py @@ -15,7 +15,8 @@ import math from dataclasses import dataclass -from typing import List, Optional, Tuple +from numbers import Integral +from typing import List, Optional, Sequence, Tuple, Union import torch import torch.nn as nn @@ -46,6 +47,7 @@ class StreamingSortformerState: spk_perm (torch.Tensor): Speaker permutation information for the speaker cache mean_sil_emb (torch.Tensor): Mean silence embedding n_sil_frames (torch.Tensor): Number of silence frames + max_speakers (torch.Tensor): Per-row number of enabled speaker channels, or ``None`` for all channels. """ spkcache = None # Speaker cache to store embeddings from start @@ -58,6 +60,7 @@ class StreamingSortformerState: spk_perm = None mean_sil_emb = None n_sil_frames = None + max_speakers = None def to(self, device): """ @@ -86,6 +89,8 @@ def to(self, device): self.mean_sil_emb = self.mean_sil_emb.to(device) if self.n_sil_frames is not None: self.n_sil_frames = self.n_sil_frames.to(device) + if self.max_speakers is not None: + self.max_speakers = self.max_speakers.to(device) class SortformerModules(NeuralModule, Exportable): @@ -486,7 +491,13 @@ def concat_and_pad(embs: List[torch.Tensor], lengths: List[torch.Tensor], output output = flat_output[: batch_size * sig_length].view(batch_size, sig_length, emb_dim) return output, total_lengths - def init_streaming_state(self, batch_size: int = 1, async_streaming: bool = False, device: torch.device = None): + def init_streaming_state( + self, + batch_size: int = 1, + async_streaming: bool = False, + device: torch.device = None, + max_speakers: Optional[Union[int, Sequence[int], torch.Tensor]] = None, + ): """ Initializes StreamingSortformerState with empty tensors or zero-valued tensors. @@ -494,6 +505,8 @@ def init_streaming_state(self, batch_size: int = 1, async_streaming: bool = Fals batch_size (int): Batch size for tensors in streaming state async_streaming (bool): True for asynchronous update, False for synchronous update device (torch.device): Device for tensors in streaming state + max_speakers (Optional[Union[int, Sequence[int], torch.Tensor]]): Number of enabled speaker channels for + every row. A scalar applies to the complete batch; a sequence or tensor supplies one value per row. Returns: streaming_state (SortformerStreamingState): initialized streaming state @@ -511,8 +524,51 @@ def init_streaming_state(self, batch_size: int = 1, async_streaming: bool = Fals streaming_state.fifo = torch.zeros((batch_size, 0, self.fc_d_model), device=device) streaming_state.mean_sil_emb = torch.zeros((batch_size, self.fc_d_model), device=device) streaming_state.n_sil_frames = torch.zeros((batch_size,), dtype=torch.long, device=device) + streaming_state.max_speakers = self._normalize_max_speakers(max_speakers, batch_size, device) return streaming_state + def _normalize_max_speakers(self, max_speakers, batch_size, device): + """Validate a scalar or per-row speaker limit and return it as a device tensor.""" + if max_speakers is None: + return None + if isinstance(max_speakers, bool): + raise TypeError("max_speakers must be an integer, sequence of integers, or integer tensor") + if isinstance(max_speakers, Integral): + max_speakers = torch.full((batch_size,), int(max_speakers), dtype=torch.long, device=device) + elif isinstance(max_speakers, torch.Tensor): + if ( + max_speakers.dtype == torch.bool + or torch.is_floating_point(max_speakers) + or torch.is_complex(max_speakers) + ): + raise TypeError("max_speakers must contain integers") + if max_speakers.ndim == 0: + max_speakers = max_speakers.expand(batch_size) + elif max_speakers.ndim != 1 or max_speakers.numel() != batch_size: + raise ValueError(f"max_speakers must contain one value per batch row; expected {batch_size}") + max_speakers = max_speakers.to(device=device, dtype=torch.long) + else: + if not isinstance(max_speakers, Sequence) or isinstance(max_speakers, (str, bytes)): + raise TypeError("max_speakers must be an integer, sequence of integers, or integer tensor") + if len(max_speakers) != batch_size: + raise ValueError(f"max_speakers must contain one value per batch row; expected {batch_size}") + if any(isinstance(value, bool) or not isinstance(value, Integral) for value in max_speakers): + raise TypeError("max_speakers must contain integers") + max_speakers = torch.tensor(max_speakers, dtype=torch.long, device=device) + + if torch.any(max_speakers < 1) or torch.any(max_speakers > self.n_spk): + raise ValueError(f"max_speakers values must be between 1 and {self.n_spk}") + return max_speakers + + @staticmethod + def apply_max_speakers_mask(predictions, max_speakers): + """Zero speaker channels at or above each row's configured limit.""" + if max_speakers is None: + return predictions + speaker_indices = torch.arange(predictions.shape[2], device=predictions.device).view(1, 1, -1) + enabled_speakers = speaker_indices < max_speakers.view(-1, 1, 1) + return predictions.masked_fill(~enabled_speakers, 0.0) + @staticmethod def apply_mask_to_preds(spkcache_fifo_chunk_preds, spkcache_fifo_chunk_fc_encoder_lengths): """ @@ -647,6 +703,8 @@ def _update_async_fifo( max_chunk_len, max_pop_out_len, lc, + silence_fifo_preds=None, + silence_chunk_preds=None, ): """ Pop and retain logical ``[FIFO | chunk]`` frames and mutate the FIFO streaming state. @@ -668,12 +726,19 @@ def _update_async_fifo( max_chunk_len (int): Physical chunk capacity excluding context. max_pop_out_len (int): Physical capacity of the rectangular popped-frame buffer. lc (int): Left context offset of the current chunk. + silence_fifo_preds (Optional[torch.Tensor]): Unmasked predictions for the valid current FIFO region, used + only to decide whether popped frames are silent. + silence_chunk_preds (Optional[torch.Tensor]): Unmasked predictions for the current chunk region, used only + to decide whether popped frames are silent. Returns: pop_out_embs (torch.Tensor): Left-aligned popped embeddings. Shape: (batch_size, max_pop_out_len, emb_dim) pop_out_preds (torch.Tensor): Left-aligned predictions for popped embeddings. Shape: (batch_size, max_pop_out_len, n_spk) + pop_out_silence_preds (torch.Tensor): Predictions used only for silence classification. This is identical + to ``pop_out_preds`` unless unmasked silence-profile predictions were supplied. + Shape: (batch_size, max_pop_out_len, n_spk) valid_pop_mask (torch.Tensor): Mask identifying valid popped frames. Shape: (batch_size, max_pop_out_len) """ @@ -726,12 +791,27 @@ def _update_async_fifo( ) pop_out_embs, updated_fifo = torch.split(gathered_fifo_embs, [max_pop_out_len, max_fifo_len], dim=1) pop_out_preds, updated_fifo_preds = torch.split(gathered_fifo_preds, [max_pop_out_len, max_fifo_len], dim=1) + pop_out_silence_preds = pop_out_preds + if silence_fifo_preds is not None and silence_chunk_preds is not None: + silence_fifo_chunk_preds = torch.cat( + [ + silence_fifo_preds, + silence_chunk_preds, + silence_fifo_preds.new_zeros((batch_size, 1, n_spk)), + ], + dim=1, + ) + pop_out_silence_preds = torch.gather( + silence_fifo_chunk_preds, + 1, + fifo_physical_indices[:, :max_pop_out_len].unsqueeze(-1).expand(-1, -1, n_spk), + ) valid_pop_mask = pop_positions < pop_out_lengths.unsqueeze(1) streaming_state.fifo = updated_fifo streaming_state.fifo_preds = updated_fifo_preds streaming_state.fifo_lengths.copy_(new_fifo_lengths) - return pop_out_embs, pop_out_preds, valid_pop_mask + return pop_out_embs, pop_out_preds, pop_out_silence_preds, valid_pop_mask def _update_async_silence_profile(self, streaming_state, pop_out_embs, pop_out_preds, valid_pop_mask): """ @@ -846,7 +926,16 @@ def _update_async_spkcache( streaming_state.spkcache_compressed[idx] = True streaming_state.spkcache_lengths.copy_(updated_spkcache_lengths.clamp(max=self.spkcache_len)) - def streaming_update_async(self, streaming_state, chunk, chunk_lengths, preds, lc: int = 0, rc: int = 0): + def streaming_update_async( + self, + streaming_state, + chunk, + chunk_lengths, + preds, + lc: int = 0, + rc: int = 0, + silence_profile_preds=None, + ): """ Update the speaker cache and FIFO queue with the chunk of embeddings and speaker predictions. Asynchronous version, which means speaker cache, FIFO and chunk may have different lengths within a batch. @@ -862,6 +951,8 @@ def streaming_update_async(self, streaming_state, chunk, chunk_lengths, preds, l Shape: (batch_size, spkcache_len + fifo_len + lc+chunk_len+rc, num_spks) lc (int): Left-context offset. Only ``chunk[:, lc:chunk_len+lc]`` is used to update the state. rc (int): Right-context offset excluded from the speaker-cache and FIFO update. + silence_profile_preds (Optional[torch.Tensor]): Unmasked speaker predictions used only to classify popped + frames as speech or silence. State and returned predictions continue to use ``preds``. Returns: streaming_state (SortformerStreamingState): Current streaming state including speaker cache and FIFO @@ -893,11 +984,25 @@ def streaming_update_async(self, streaming_state, chunk, chunk_lengths, preds, l max_chunk_len, lc, ) + silence_fifo_preds = None + silence_chunk_preds = None + if silence_profile_preds is not None: + _, silence_fifo_preds, silence_chunk_preds = self._gather_async_predictions( + streaming_state, + silence_profile_preds, + spkcache_lengths, + fifo_lengths, + chunk_lengths, + max_spkcache_len, + max_fifo_len, + max_chunk_len, + lc, + ) pop_out_lengths, new_fifo_lengths = self._compute_async_fifo_pop_lengths( spkcache_lengths, fifo_lengths, chunk_lengths, max_fifo_len ) - pop_out_embs, pop_out_preds, valid_pop_mask = self._update_async_fifo( + pop_out_embs, pop_out_preds, pop_out_silence_preds, valid_pop_mask = self._update_async_fifo( streaming_state, chunk, current_fifo_preds, @@ -908,8 +1013,10 @@ def streaming_update_async(self, streaming_state, chunk, chunk_lengths, preds, l max_chunk_len, max_pop_out_len, lc, + silence_fifo_preds=silence_fifo_preds, + silence_chunk_preds=silence_chunk_preds, ) - self._update_async_silence_profile(streaming_state, pop_out_embs, pop_out_preds, valid_pop_mask) + self._update_async_silence_profile(streaming_state, pop_out_embs, pop_out_silence_preds, valid_pop_mask) self._update_async_spkcache( streaming_state, current_spkcache_preds, @@ -928,7 +1035,15 @@ def streaming_update_async(self, streaming_state, chunk, chunk_lengths, preds, l return streaming_state, chunk_preds - def streaming_update(self, streaming_state, chunk, preds, lc: int = 0, rc: int = 0): + def streaming_update( + self, + streaming_state, + chunk, + preds, + lc: int = 0, + rc: int = 0, + silence_profile_preds=None, + ): """ Update the speaker cache and FIFO queue with the chunk of embeddings and speaker predictions. Synchronous version, which means speaker cahce, FIFO queue and chunk have same lengths within a batch. @@ -942,6 +1057,8 @@ def streaming_update(self, streaming_state, chunk, preds, lc: int = 0, rc: int = Shape: (batch_size, spkcache_len + fifo_len + lc+chunk_len+rc, num_spks) lc (int): Left-context offset. Only ``chunk[:, lc:chunk_len+lc]`` is used to update the state. rc (int): Right-context offset excluded from the speaker-cache and FIFO update. + silence_profile_preds (Optional[torch.Tensor]): Unmasked speaker predictions used only to classify popped + frames as speech or silence. State and returned predictions continue to use ``preds``. Returns: streaming_state (SortformerStreamingState): current streaming state including speaker cache and FIFO @@ -963,10 +1080,26 @@ def streaming_update(self, streaming_state, chunk, preds, lc: int = 0, rc: int = preds = torch.stack( [preds[batch_index, :, inv_spk_perm[batch_index]] for batch_index in range(batch_size)] ) + if silence_profile_preds is not None: + silence_profile_preds = torch.stack( + [ + silence_profile_preds[batch_index, :, inv_spk_perm[batch_index]] + for batch_index in range(batch_size) + ] + ) streaming_state.fifo_preds = preds[:, spkcache_len : spkcache_len + fifo_len] chunk = chunk[:, lc : chunk_len + lc] chunk_preds = preds[:, spkcache_len + fifo_len + lc : spkcache_len + fifo_len + chunk_len + lc] + silence_profile_fifo_preds = None + if silence_profile_preds is not None: + silence_profile_fifo_preds = torch.cat( + [ + silence_profile_preds[:, spkcache_len : spkcache_len + fifo_len], + silence_profile_preds[:, spkcache_len + fifo_len + lc : spkcache_len + fifo_len + chunk_len + lc], + ], + dim=1, + ) # append chunk to fifo streaming_state.fifo = torch.cat([streaming_state.fifo, chunk], dim=1) @@ -980,12 +1113,15 @@ def streaming_update(self, streaming_state, chunk, preds, lc: int = 0, rc: int = pop_out_embs = streaming_state.fifo[:, :pop_out_len] pop_out_preds = streaming_state.fifo_preds[:, :pop_out_len] + silence_profile_pop_out_preds = ( + pop_out_preds if silence_profile_fifo_preds is None else silence_profile_fifo_preds[:, :pop_out_len] + ) if not self.use_learnable_sil_emb: streaming_state.mean_sil_emb, streaming_state.n_sil_frames = self._get_silence_profile( streaming_state.mean_sil_emb, streaming_state.n_sil_frames, pop_out_embs, - pop_out_preds, + silence_profile_pop_out_preds, ) streaming_state.fifo = streaming_state.fifo[:, pop_out_len:] streaming_state.fifo_preds = streaming_state.fifo_preds[:, pop_out_len:] diff --git a/nemo/collections/asr/parts/utils/sortformer_utils.py b/nemo/collections/asr/parts/utils/sortformer_utils.py index 0b6657f9f619..3fd96a92d604 100644 --- a/nemo/collections/asr/parts/utils/sortformer_utils.py +++ b/nemo/collections/asr/parts/utils/sortformer_utils.py @@ -13,21 +13,314 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import logging +import math import os import time from functools import wraps from pathlib import Path from tempfile import NamedTemporaryFile -from typing import TYPE_CHECKING, Dict, List, Optional +from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union import torch from omegaconf import open_dict +from nemo.collections.asr.parts.preprocessing.features import normalize_batch + if TYPE_CHECKING: from nemo.collections.asr.models import SortformerEncLabelModel +class SortformerStreamingSession: + """Stateful raw-audio session for a fixed batch of streaming Sortformer inputs. + + Each row owns independent waveform buffering, progress, finalization, and speaker-cache state. Incoming waveform + chunks may have different valid lengths, and the session uses Sortformer's asynchronous streaming update so idle, + active, and finalized rows can coexist in one batch. + + Args: + model: Streaming ``SortformerEncLabelModel`` in evaluation mode. + batch_size: Fixed number of independent audio streams owned by the session. + max_speakers: Number of enabled speaker channels for every stream. A scalar applies to the complete batch; a + sequence or tensor supplies one value per row. By default, every model speaker channel is enabled. + """ + + def __init__( + self, + model: "SortformerEncLabelModel", + batch_size: int = 1, + max_speakers: Optional[Union[int, Sequence[int], torch.Tensor]] = None, + ): + if not model.streaming_mode: + raise ValueError("SortformerStreamingSession requires a model with streaming_mode=True") + if model.training: + raise ValueError("SortformerStreamingSession requires an evaluation model; call model.eval() first") + if not isinstance(batch_size, int) or isinstance(batch_size, bool) or batch_size < 1: + raise ValueError(f"batch_size must be a positive integer, got {batch_size}") + + self.model = model + self.batch_size = batch_size + self._max_speakers = max_speakers + self.device = model.device + self._normalization = model.preprocessor.featurizer.normalize + self._preprocessor = copy.deepcopy(model.preprocessor).to(self.device).eval() + self._preprocessor.featurizer.normalize = None + self._preprocessor.featurizer.dither = 0.0 + self._preprocessor.featurizer.pad_to = 0 + + self._hop_length = self._preprocessor.hop_length + self._n_fft = self._preprocessor.featurizer.n_fft + self._stft_margin_frames = math.ceil((self._n_fft // 2 + 1) / self._hop_length) + 1 + self._chunk_frames = model.sortformer_modules.chunk_len * model.encoder.subsampling_factor + self._left_context_frames = model.sortformer_modules.chunk_left_context * model.encoder.subsampling_factor + self._right_context_frames = model.sortformer_modules.chunk_right_context * model.encoder.subsampling_factor + self.reset() + + @torch.inference_mode() + def diarize_step( + self, + audio_chunks: torch.Tensor, + audio_chunk_lengths: Optional[torch.Tensor] = None, + is_final: Union[bool, torch.Tensor] = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Consume one raw-audio chunk per stream and return newly committed speaker probabilities. + + Args: + audio_chunks: Float waveforms with shape ``(batch_size, max_num_samples)``. A one-dimensional tensor is + also accepted when ``batch_size=1``. + audio_chunk_lengths: Valid sample count for each padded row, with shape ``(batch_size,)``. If omitted, + every row uses the complete waveform width. + is_final: Boolean finalization mask with shape ``(batch_size,)`` or one boolean applied to every row. + Finalized rows can remain in later calls with a zero audio length while other rows continue. + + Returns: + padded_probabilities: Newly committed probabilities with shape + ``(batch_size, max_new_frames, num_speakers)``. + probability_lengths: Valid output frames for each row, with shape ``(batch_size,)``. + """ + audio_chunks, audio_chunk_lengths = self._validate_audio_chunks(audio_chunks, audio_chunk_lengths) + final_mask = self._validate_final_mask(is_final) + input_lengths = audio_chunk_lengths.tolist() + final_flags = final_mask.tolist() + + for stream_index, chunk_length in enumerate(input_lengths): + if self._finalized[stream_index] and chunk_length > 0: + raise RuntimeError( + f"Cannot supply audio to finalized stream {stream_index}; call reset() before reusing the session" + ) + if chunk_length > 0: + chunk = audio_chunks[stream_index, :chunk_length] + self._audio_buffers[stream_index] = torch.cat([self._audio_buffers[stream_index], chunk]) + self._received_samples[stream_index] += chunk_length + + emitted = [[] for _ in range(self.batch_size)] + while True: + ready_groups = self._get_ready_groups(final_flags) + if not ready_groups: + break + + for (left_offset, right_offset), requests in ready_groups.items(): + processed_signal, processed_signal_length = self._extract_feature_batch(requests) + empty_preds = processed_signal.new_zeros((self.batch_size, 0, self.model.sortformer_modules.n_spk)) + self.streaming_state, chunk_preds = self.model.forward_streaming_step( + processed_signal=processed_signal, + processed_signal_length=processed_signal_length, + streaming_state=self.streaming_state, + total_preds=empty_preds, + left_offset=left_offset, + right_offset=right_offset, + async_streaming=True, + ) + + for stream_index, _, _, central_end in requests: + committed_feature_frames = central_end - self._next_feature_frames[stream_index] + output_length = math.ceil(committed_feature_frames / self.model.output_subsampling_factor) + if output_length > chunk_preds.shape[1]: + raise RuntimeError( + "Streaming model returned fewer prediction frames than required: " + f"needed {output_length}, got {chunk_preds.shape[1]}" + ) + emitted[stream_index].append(chunk_preds[stream_index, :output_length]) + self._next_feature_frames[stream_index] = central_end + + self._compact_audio_buffers() + for stream_index, is_final_stream in enumerate(final_flags): + if is_final_stream: + self._finalized[stream_index] = True + + return self._pad_emitted_outputs(emitted) + + def reset(self) -> None: + """Clear every stream's buffered audio and initialize a fresh batched asynchronous model state.""" + self.streaming_state = self.model.sortformer_modules.init_streaming_state( + batch_size=self.batch_size, + async_streaming=True, + device=self.device, + max_speakers=self._max_speakers, + ) + self._max_speakers = self.streaming_state.max_speakers + self._audio_buffers = [torch.empty(0, dtype=torch.float32, device=self.device) for _ in range(self.batch_size)] + self._audio_buffer_starts = [0] * self.batch_size + self._received_samples = [0] * self.batch_size + self._next_feature_frames = [0] * self.batch_size + self._finalized = [False] * self.batch_size + + def _get_ready_groups(self, final_flags: List[bool]): + ready_groups = {} + for stream_index in range(self.batch_size): + if self._finalized[stream_index]: + continue + available_frames = self._available_feature_frames(stream_index, is_final=final_flags[stream_index]) + next_frame = self._next_feature_frames[stream_index] + if next_frame >= available_frames: + continue + + central_end = min(next_frame + self._chunk_frames, available_frames) + if not final_flags[stream_index] and central_end + self._right_context_frames > available_frames: + continue + + feature_start = max(0, next_frame - self._left_context_frames) + feature_end = min(central_end + self._right_context_frames, available_frames) + offsets = (next_frame - feature_start, feature_end - central_end) + ready_groups.setdefault(offsets, []).append((stream_index, feature_start, feature_end, central_end)) + return ready_groups + + def _available_feature_frames(self, stream_index: int, is_final: bool) -> int: + received_samples = self._received_samples[stream_index] + sample_count = torch.tensor(received_samples, device=self.device) + offline_frames = int(self._preprocessor.featurizer.get_seq_len(sample_count).item()) + if is_final: + return max(0, offline_frames) + + stable_samples = received_samples - self._n_fft // 2 + if stable_samples < 0: + return 0 + stable_frames = stable_samples // self._hop_length + 1 + return max(0, min(offline_frames, stable_frames)) + + def _compact_audio_buffers(self) -> None: + for stream_index in range(self.batch_size): + first_needed_frame = max( + 0, + self._next_feature_frames[stream_index] - self._left_context_frames - self._stft_margin_frames, + ) + first_needed_sample = first_needed_frame * self._hop_length + drop_samples = first_needed_sample - self._audio_buffer_starts[stream_index] + if drop_samples > 0: + self._audio_buffers[stream_index] = self._audio_buffers[stream_index][drop_samples:].clone() + self._audio_buffer_starts[stream_index] = first_needed_sample + + def _extract_feature_batch(self, requests): + audio_segments = [] + audio_lengths = [] + local_feature_ranges = [] + for stream_index, feature_start, feature_end, _ in requests: + segment_start_frame = max(0, feature_start - self._stft_margin_frames) + segment_start_sample = segment_start_frame * self._hop_length + segment_end_sample = min( + self._received_samples[stream_index], + (feature_end + self._stft_margin_frames) * self._hop_length, + ) + buffer_start = segment_start_sample - self._audio_buffer_starts[stream_index] + buffer_end = segment_end_sample - self._audio_buffer_starts[stream_index] + audio_segment = self._audio_buffers[stream_index][buffer_start:buffer_end] + audio_segments.append(audio_segment) + audio_lengths.append(audio_segment.shape[0]) + local_start = feature_start - segment_start_frame + local_feature_ranges.append((local_start, local_start + feature_end - feature_start)) + + padded_audio = torch.nn.utils.rnn.pad_sequence(audio_segments, batch_first=True) + audio_lengths = torch.tensor(audio_lengths, dtype=torch.long, device=self.device) + features, feature_lengths = self._preprocessor(input_signal=padded_audio, length=audio_lengths) + + feature_windows = [] + window_lengths = [] + for request_index, (local_start, local_end) in enumerate(local_feature_ranges): + if feature_lengths[request_index] < local_end: + raise RuntimeError( + "Streaming preprocessor returned fewer feature frames than required: " + f"needed {local_end}, got {feature_lengths[request_index].item()}" + ) + window = features[request_index, :, local_start:local_end].transpose(0, 1) + feature_windows.append(window) + window_lengths.append(window.shape[0]) + + active_features = torch.nn.utils.rnn.pad_sequence(feature_windows, batch_first=True).transpose(1, 2) + active_lengths = torch.tensor(window_lengths, dtype=torch.long, device=self.device) + if self._normalization: + active_features, _, _ = normalize_batch(active_features, active_lengths, self._normalization) + active_features = active_features.transpose(1, 2) + + batch_features = active_features.new_zeros( + (self.batch_size, active_features.shape[1], active_features.shape[2]) + ) + batch_lengths = torch.zeros(self.batch_size, dtype=torch.long, device=self.device) + for request_index, (stream_index, _, _, _) in enumerate(requests): + feature_length = window_lengths[request_index] + batch_features[stream_index, :feature_length] = active_features[request_index, :feature_length] + batch_lengths[stream_index] = feature_length + return batch_features, batch_lengths + + def _pad_emitted_outputs(self, emitted): + num_speakers = self.model.sortformer_modules.n_spk + dtype = next(self.model.parameters()).dtype + row_outputs = [ + torch.cat(row, dim=0) if row else torch.zeros((0, num_speakers), dtype=dtype, device=self.device) + for row in emitted + ] + output_lengths = torch.tensor([row.shape[0] for row in row_outputs], dtype=torch.long, device=self.device) + max_output_length = max((row.shape[0] for row in row_outputs), default=0) + padded_outputs = torch.zeros( + (self.batch_size, max_output_length, num_speakers), dtype=dtype, device=self.device + ) + for stream_index, row in enumerate(row_outputs): + padded_outputs[stream_index, : row.shape[0]] = row + return padded_outputs, output_lengths + + def _validate_audio_chunks(self, audio_chunks, audio_chunk_lengths): + if not isinstance(audio_chunks, torch.Tensor): + raise TypeError(f"audio_chunks must be a torch.Tensor, got {type(audio_chunks).__name__}") + if audio_chunks.ndim == 1 and self.batch_size == 1: + audio_chunks = audio_chunks.unsqueeze(0) + if audio_chunks.ndim != 2 or audio_chunks.shape[0] != self.batch_size: + raise ValueError( + f"audio_chunks must have batch dimension {self.batch_size} and shape " + f"({self.batch_size}, max_num_samples); got {tuple(audio_chunks.shape)}" + ) + audio_chunks = audio_chunks.detach().to(device=self.device, dtype=torch.float32) + + if audio_chunk_lengths is None: + audio_chunk_lengths = torch.full( + (self.batch_size,), audio_chunks.shape[1], dtype=torch.long, device=self.device + ) + elif not isinstance(audio_chunk_lengths, torch.Tensor): + raise TypeError( + f"audio_chunk_lengths must be a torch.Tensor or None, got {type(audio_chunk_lengths).__name__}" + ) + elif audio_chunk_lengths.shape != (self.batch_size,): + raise ValueError( + f"audio_chunk_lengths must have shape ({self.batch_size},), got {tuple(audio_chunk_lengths.shape)}" + ) + elif torch.is_floating_point(audio_chunk_lengths) or audio_chunk_lengths.dtype == torch.bool: + raise TypeError("audio_chunk_lengths must contain integers") + else: + audio_chunk_lengths = audio_chunk_lengths.to(device=self.device, dtype=torch.long) + + if torch.any(audio_chunk_lengths < 0) or torch.any(audio_chunk_lengths > audio_chunks.shape[1]): + raise ValueError(f"audio_chunk_lengths must be between 0 and {audio_chunks.shape[1]} samples") + return audio_chunks, audio_chunk_lengths + + def _validate_final_mask(self, is_final): + if isinstance(is_final, bool): + return torch.full((self.batch_size,), is_final, dtype=torch.bool, device=self.device) + if not isinstance(is_final, torch.Tensor): + raise TypeError(f"is_final must be a boolean or torch.Tensor, got {type(is_final).__name__}") + if is_final.dtype != torch.bool or is_final.shape != (self.batch_size,): + raise ValueError(f"is_final must be boolean with shape ({self.batch_size},), got {tuple(is_final.shape)}") + return is_final.to(device=self.device) + + def configure_output_subsampling_factor( diar_model: "SortformerEncLabelModel", output_subsampling_factor: Optional[int], diff --git a/tests/collections/speaker_tasks/test_diar_sortformer_models.py b/tests/collections/speaker_tasks/test_diar_sortformer_models.py index 1ee476221d49..b1163c03d33d 100644 --- a/tests/collections/speaker_tasks/test_diar_sortformer_models.py +++ b/tests/collections/speaker_tasks/test_diar_sortformer_models.py @@ -268,6 +268,169 @@ def test_constructor(self, sortformer_model): instance2 = SortformerEncLabelModel.from_config_dict(confdict) assert isinstance(instance2, SortformerEncLabelModel) + @pytest.mark.unit + def test_raw_audio_streaming_session_batches_independent_streams(self): + model = _create_sortformer_model().eval() + model.streaming_mode = True + model.sortformer_modules.chunk_len = 2 + model.sortformer_modules.chunk_left_context = 1 + model.sortformer_modules.chunk_right_context = 1 + model._check_streaming_parameters() + audio = [torch.randn(8193), torch.randn(5001)] + + reference_preds = [] + for signal in audio: + session = model.create_streaming_session(batch_size=1) + preds, pred_lengths = session.diarize_step( + signal.unsqueeze(0), + audio_chunk_lengths=torch.tensor([signal.numel()]), + is_final=torch.tensor([True]), + ) + reference_preds.append(preds[0, : pred_lengths[0]]) + + session = model.create_streaming_session(batch_size=2) + emitted = [[], []] + offsets = [0, 0] + step_sizes = [(17, 503), (1600, 81), (2999, 4417), (3577, 0)] + for step_index, sizes in enumerate(step_sizes): + chunks = [] + lengths = [] + final = [] + for stream_index, size in enumerate(sizes): + end = min(offsets[stream_index] + size, audio[stream_index].numel()) + chunks.append(audio[stream_index][offsets[stream_index] : end]) + offsets[stream_index] = end + lengths.append(chunks[-1].numel()) + final.append(end == audio[stream_index].numel()) + padded_audio = torch.nn.utils.rnn.pad_sequence(chunks, batch_first=True) + preds, pred_lengths = session.diarize_step( + padded_audio, + audio_chunk_lengths=torch.tensor(lengths), + is_final=torch.tensor(final), + ) + for stream_index in range(2): + emitted[stream_index].append(preds[stream_index, : pred_lengths[stream_index]]) + + if step_index == 2: + assert final == [False, True] + + for stream_index in range(2): + batched_preds = torch.cat(emitted[stream_index]) + torch.testing.assert_close(batched_preds, reference_preds[stream_index]) + assert ( + session._audio_buffers[stream_index].untyped_storage().nbytes() + < audio[stream_index].untyped_storage().nbytes() + ) + + @pytest.mark.unit + @pytest.mark.parametrize("high_resolution", [False, True]) + def test_raw_audio_streaming_session_applies_per_row_max_speakers_before_state_update(self, high_resolution): + model = _create_sortformer_model(high_resolution=high_resolution).eval() + model.streaming_mode = True + model.sortformer_modules.chunk_len = 2 + model.sortformer_modules.chunk_left_context = 1 + model.sortformer_modules.chunk_right_context = 1 + model._check_streaming_parameters() + captured_state_predictions = [] + captured_silence_profile_predictions = [] + streaming_update_async = model.sortformer_modules.streaming_update_async + + def capture_streaming_update(**kwargs): + captured_state_predictions.append(kwargs["preds"].clone()) + captured_silence_profile_predictions.append(kwargs["silence_profile_preds"].clone()) + return streaming_update_async(**kwargs) + + model.sortformer_modules.streaming_update_async = capture_streaming_update + session = model.create_streaming_session(batch_size=2, max_speakers=[2, 4]) + audio = torch.nn.utils.rnn.pad_sequence([torch.randn(8193), torch.randn(5001)], batch_first=True) + predictions, prediction_lengths = session.diarize_step( + audio, + audio_chunk_lengths=torch.tensor([8193, 5001]), + is_final=torch.tensor([True, True]), + ) + + assert torch.count_nonzero(predictions[0, : prediction_lengths[0], 2:]) == 0 + assert torch.count_nonzero(predictions[1, : prediction_lengths[1], 2:]) > 0 + assert captured_state_predictions + for state_predictions in captured_state_predictions: + assert torch.count_nonzero(state_predictions[0, :, 2:]) == 0 + assert any( + torch.count_nonzero(silence_predictions[0, :, 2:]) > 0 + for silence_predictions in captured_silence_profile_predictions + ) + assert any( + torch.count_nonzero(state_predictions[1, :, 2:]) > 0 for state_predictions in captured_state_predictions + ) + assert torch.count_nonzero(session.streaming_state.fifo_preds[0, :, 2:]) == 0 + assert torch.count_nonzero(session.streaming_state.spkcache_preds[0, :, 2:]) == 0 + assert session.streaming_state.max_speakers.tolist() == [2, 4] + session.reset() + assert session.streaming_state.max_speakers.tolist() == [2, 4] + + @pytest.mark.unit + @pytest.mark.parametrize( + "max_speakers, error_type, error_match", + [ + ([2], ValueError, "one value per batch row"), + ([2, 5], ValueError, "between 1 and 4"), + ([2, 1.5], TypeError, "must contain integers"), + (torch.tensor([2.0, 4.0]), TypeError, "must contain integers"), + (True, TypeError, "must be an integer"), + ], + ) + def test_raw_audio_streaming_session_rejects_invalid_max_speakers(self, max_speakers, error_type, error_match): + model = _create_sortformer_model().eval() + model.streaming_mode = True + + with pytest.raises(error_type, match=error_match): + model.create_streaming_session(batch_size=2, max_speakers=max_speakers) + + @pytest.mark.unit + @pytest.mark.parametrize( + "max_speakers, expected", + [ + (2, [2, 2]), + (torch.tensor(3), [3, 3]), + (torch.tensor([1, 4]), [1, 4]), + ], + ) + def test_raw_audio_streaming_session_normalizes_max_speakers(self, max_speakers, expected): + model = _create_sortformer_model().eval() + model.streaming_mode = True + + session = model.create_streaming_session(batch_size=2, max_speakers=max_speakers) + + assert session.streaming_state.max_speakers.tolist() == expected + + @pytest.mark.unit + def test_raw_audio_streaming_session_reset_and_validation(self): + offline_model = _create_sortformer_model().eval() + with pytest.raises(ValueError, match="streaming_mode=True"): + offline_model.create_streaming_session() + + model = _create_sortformer_model().eval() + model.streaming_mode = True + model.sortformer_modules.chunk_len = 2 + model.sortformer_modules.chunk_left_context = 1 + model.sortformer_modules.chunk_right_context = 1 + model._check_streaming_parameters() + audio = torch.randn(4097) + session = model.create_streaming_session(batch_size=1) + + first_preds, first_lengths = session.diarize_step(audio, is_final=True) + with pytest.raises(RuntimeError, match="finalized stream 0"): + session.diarize_step(torch.ones(1)) + session.reset() + second_preds, second_lengths = session.diarize_step(audio.unsqueeze(0), is_final=torch.tensor([True])) + + assert torch.equal(second_lengths, first_lengths) + torch.testing.assert_close(second_preds, first_preds) + with pytest.raises(ValueError, match="batch dimension"): + session.reset() + session.diarize_step(torch.randn(2, 100)) + with pytest.raises(ValueError, match="positive integer"): + model.create_streaming_session(batch_size=0) + @pytest.mark.unit @pytest.mark.parametrize( "stacking_factor, feature_shape, input_lengths, expected_encoded_lengths", diff --git a/tests/collections/speaker_tasks/test_diar_sortformer_modules.py b/tests/collections/speaker_tasks/test_diar_sortformer_modules.py index 56e7527788d4..9f646cb4f365 100644 --- a/tests/collections/speaker_tasks/test_diar_sortformer_modules.py +++ b/tests/collections/speaker_tasks/test_diar_sortformer_modules.py @@ -2513,6 +2513,48 @@ def test_async_zero_capacity_fifo_masks_ragged_silence_updates( assert streaming_state.n_sil_frames.tolist() == list(expected_silence_counts) torch.testing.assert_close(streaming_state.mean_sil_emb, torch.tensor(expected_silence_means)) + @pytest.mark.unit + @pytest.mark.parametrize("async_streaming", [False, True]) + def test_speaker_limit_does_not_turn_disabled_channel_activity_into_silence(self, async_streaming): + sortformer_modules = SortformerModules( + num_spks=4, + fc_d_model=2, + spkcache_len=8, + fifo_len=0, + chunk_len=2, + spkcache_update_period=2, + spkcache_sil_frames_per_spk=0, + use_learnable_sil_emb=False, + ) + streaming_state = sortformer_modules.init_streaming_state( + batch_size=1, async_streaming=async_streaming, max_speakers=2 + ) + chunk = torch.tensor([[[1.0, 3.0], [5.0, 7.0]]]) + silence_profile_preds = torch.zeros(1, 2, 4) + silence_profile_preds[0, 0, 2] = 1.0 + state_preds = sortformer_modules.apply_max_speakers_mask(silence_profile_preds, streaming_state.max_speakers) + + if async_streaming: + streaming_state, chunk_preds = sortformer_modules.streaming_update_async( + streaming_state, + chunk, + torch.tensor([2]), + state_preds, + silence_profile_preds=silence_profile_preds, + ) + else: + streaming_state, chunk_preds = sortformer_modules.streaming_update( + streaming_state, + chunk, + state_preds, + silence_profile_preds=silence_profile_preds, + ) + + assert streaming_state.n_sil_frames.tolist() == [1] + torch.testing.assert_close(streaming_state.mean_sil_emb, chunk[:, 1]) + assert torch.count_nonzero(streaming_state.spkcache_preds[..., 2:]) == 0 + assert torch.count_nonzero(chunk_preds[..., 2:]) == 0 + @pytest.mark.unit @pytest.mark.parametrize( (