Add per-stream speaker limits to Sortformer sessions

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 |

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