WhisperLiveKit/whisperlivekit/diarization/sortformer_backend_offline.py

import numpy as np
import torch
import logging

from nemo.collections.asr.models import SortformerEncLabelModel
from nemo.collections.asr.modules import AudioToMelSpectrogramPreprocessor
import librosa

logger = logging.getLogger(__name__)

def load_model():

    diar_model = SortformerEncLabelModel.from_pretrained("nvidia/diar_streaming_sortformer_4spk-v2")
    diar_model.eval()

    if torch.cuda.is_available():
        diar_model.to(torch.device("cuda"))

    #we target 1 second lag for the moment. chunk_len could be reduced.
    diar_model.sortformer_modules.chunk_len = 10
    diar_model.sortformer_modules.subsampling_factor = 10 #8 would be better ideally

    diar_model.sortformer_modules.chunk_right_context = 0 #no.
    diar_model.sortformer_modules.chunk_left_context = 10 #big so it compensiate the problem with no padding later.

    diar_model.sortformer_modules.spkcache_len = 188
    diar_model.sortformer_modules.fifo_len = 188
    diar_model.sortformer_modules.spkcache_update_period = 144
    diar_model.sortformer_modules.log = False
    diar_model.sortformer_modules._check_streaming_parameters()


    audio2mel = AudioToMelSpectrogramPreprocessor(
            window_size= 0.025, 
            normalize="NA",
            n_fft=512,
            features=128,
            pad_to=0) #pad_to 16 works better than 0. On test audio, we detect a third speaker for 1 second with pad_to=0. To solve that : increase left context to 10.

    return diar_model, audio2mel

diar_model, audio2mel = load_model()

class StreamingSortformerState:
    """
    This class creates a class instance that will be used to store the state of the
    streaming Sortformer model.

    Attributes:
        spkcache (torch.Tensor): Speaker cache to store embeddings from start
        spkcache_lengths (torch.Tensor): Lengths of the speaker cache
        spkcache_preds (torch.Tensor): The speaker predictions for the speaker cache parts
        fifo (torch.Tensor): FIFO queue to save the embedding from the latest chunks
        fifo_lengths (torch.Tensor): Lengths of the FIFO queue
        fifo_preds (torch.Tensor): The speaker predictions for the FIFO queue parts
        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
    """

    spkcache = None  # Speaker cache to store embeddings from start
    spkcache_lengths = None  #
    spkcache_preds = None  # speaker cache predictions
    fifo = None  # to save the embedding from the latest chunks
    fifo_lengths = None
    fifo_preds = None
    spk_perm = None
    mean_sil_emb = None
    n_sil_frames = None


def init_streaming_state(self, batch_size: int = 1, async_streaming: bool = False, device: torch.device = None):
    """
    Initializes StreamingSortformerState with empty tensors or zero-valued tensors.

    Args:
        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

    Returns:
        streaming_state (SortformerStreamingState): initialized streaming state
    """
    streaming_state = StreamingSortformerState()
    if async_streaming:
        streaming_state.spkcache = torch.zeros((batch_size, self.spkcache_len, self.fc_d_model), device=device)
        streaming_state.spkcache_preds = torch.zeros((batch_size, self.spkcache_len, self.n_spk), device=device)
        streaming_state.spkcache_lengths = torch.zeros((batch_size,), dtype=torch.long, device=device)
        streaming_state.fifo = torch.zeros((batch_size, self.fifo_len, self.fc_d_model), device=device)
        streaming_state.fifo_lengths = torch.zeros((batch_size,), dtype=torch.long, device=device)
    else:
        streaming_state.spkcache = torch.zeros((batch_size, 0, self.fc_d_model), device=device)
        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)
    return streaming_state


def process_diarization(chunks):
    """ 
    what it does:
    1. Preprocessing: Applies dithering and pre-emphasis (high-pass filter) if enabled
    2. STFT: Computes the Short-Time Fourier Transform using:
        - the window of window_size=0.025 --> size of a window : 400 samples
        - the hop parameter : n_window_stride = 0.01 -> every 160 samples, a new window
    3. Magnitude Calculation: Converts complex STFT output to magnitude spectrogram
    4. Mel Conversion: Applies Mel filterbanks (128 filters in this case) to get Mel spectrogram
    5. Logarithm: Takes the log of the Mel spectrogram (if `log=True`)
    6. Normalization: Skips normalization since `normalize="NA"`
    7. Padding: Pads the time dimension to a multiple of `pad_to` (default 16)    
    """
    previous_chunk = None
    l_chunk_feat_seq_t = []
    for chunk in chunks:
        audio_signal_chunk = torch.tensor(chunk).unsqueeze(0).to(diar_model.device)
        audio_signal_length_chunk = torch.tensor([audio_signal_chunk.shape[1]]).to(diar_model.device)
        processed_signal_chunk, processed_signal_length_chunk = audio2mel.get_features(audio_signal_chunk, audio_signal_length_chunk)
        if previous_chunk is not None:
            to_add = previous_chunk[:, :, -99:]
            total = torch.concat([to_add, processed_signal_chunk], dim=2)
        else:
            total = processed_signal_chunk
        previous_chunk = processed_signal_chunk
        l_chunk_feat_seq_t.append(torch.transpose(total, 1, 2))

    batch_size = 1
    streaming_state = init_streaming_state(diar_model.sortformer_modules,
        batch_size = batch_size,
        async_streaming = True,
        device = diar_model.device
    )
    total_preds = torch.zeros((batch_size, 0, diar_model.sortformer_modules.n_spk), device=diar_model.device)

    chunk_duration_seconds = diar_model.sortformer_modules.chunk_len * diar_model.sortformer_modules.subsampling_factor * diar_model.preprocessor._cfg.window_stride

    l_speakers = [
        {'start_time': 0,
        'end_time': 0,
        'speaker': 0
        }
    ]
    len_prediction = None
    left_offset = 0
    right_offset = 8
    for i, chunk_feat_seq_t in enumerate(l_chunk_feat_seq_t):
        with torch.inference_mode():
                streaming_state, total_preds = diar_model.forward_streaming_step(
                    processed_signal=chunk_feat_seq_t,
                    processed_signal_length=torch.tensor([chunk_feat_seq_t.shape[1]]),
                    streaming_state=streaming_state,
                    total_preds=total_preds,
                    left_offset=left_offset,
                    right_offset=right_offset,
                )
                left_offset = 8
                preds_np = total_preds[0].cpu().numpy()
                active_speakers = np.argmax(preds_np, axis=1)
                if len_prediction is None:
                    len_prediction = len(active_speakers) # we want to get the len of 1 prediction
                frame_duration = chunk_duration_seconds / len_prediction
                active_speakers = active_speakers[-len_prediction:]
                for idx, spk in enumerate(active_speakers):
                    if spk != l_speakers[-1]['speaker']:
                        l_speakers.append(
                            {'start_time': (i * chunk_duration_seconds + idx * frame_duration),
                            'end_time': (i * chunk_duration_seconds + (idx + 1) * frame_duration),
                            'speaker': spk
                        })                    
                    else:
                        l_speakers[-1]['end_time'] = i * chunk_duration_seconds + (idx + 1) * frame_duration
                    
        
        """
        Should print
        [{'start_time': 0, 'end_time': 8.72, 'speaker': 0}, 
        {'start_time': 8.72, 'end_time': 18.88, 'speaker': 1},
        {'start_time': 18.88, 'end_time': 24.96, 'speaker': 2},
        {'start_time': 24.96, 'end_time': 31.68, 'speaker': 0}]
        """
    for speaker in l_speakers:
        print(f"Speaker {speaker['speaker']}: {speaker['start_time']:.2f}s - {speaker['end_time']:.2f}s")    
    

if __name__ == '__main__':
    
    an4_audio = 'audio_test.mp3'
    signal, sr = librosa.load(an4_audio, sr=16000)
    signal = signal[:16000*30]
    # signal = signal[:-(len(signal)%16000)]

    print("\n" + "=" * 50)
    print("Expected ground truth:")
    print("Speaker 0: 0:00 - 0:09")
    print("Speaker 1: 0:09 - 0:19") 
    print("Speaker 2: 0:19 - 0:25")
    print("Speaker 0: 0:25 - 0:30")
    print("=" * 50)

    chunk_size = 16000  # 1 second
    chunks = []
    for i in range(0, len(signal), chunk_size):
        chunk = signal[i:i+chunk_size]
        chunks.append(chunk)
    
    process_diarization(chunks)