Implement new diarization model: Sortformer (#342)

* implement sortformer

* ..

* fixes

* fix audio saving

* formatting

Author: MahmoudAshraf97

.github/workflows/test_run.yml

@@ -60,6 +60,10 @@ jobs:
         python -m pip install torch torchaudio --index-url https://download.pytorch.org/whl/cpu
         uv pip install --system -c constraints.txt -r requirements.txt
 
+    - name: Test with sortformer diarizer
+      run: |
+        python diarize.py -a "./tests/assets/test.opus" --whisper-model tiny.en --diarizer sortformer
+
     - name: Test with msdd diarizer
       run: |
         python diarize.py -a "./tests/assets/test.opus" --whisper-model tiny.en --diarizer msdd

diarization/__init__.py

@@ -1,3 +1,4 @@
 from .msdd.msdd import MSDDDiarizer
+from .sortformer.sortformer import SortformerDiarizer
 
-__all__ = ["MSDDDiarizer"]
+__all__ = ["MSDDDiarizer", "SortformerDiarizer"]

diarization/sortformer/sortformer.py

@@ -0,0 +1,78 @@
+from warnings import warn
+
+import torch
+
+from nemo.collections.asr.models import SortformerEncLabelModel
+from nemo.collections.asr.parts.mixins.diarization import DiarizeConfig
+
+
+class SortformerDiarizer:
+    def __init__(self, device):
+        self.model: SortformerEncLabelModel = SortformerEncLabelModel.from_pretrained(
+            "nvidia/diar_streaming_sortformer_4spk-v2", map_location=device
+        )
+
+        self.model.sortformer_modules.chunk_len = 340
+        self.model.sortformer_modules.chunk_right_context = 40
+        self.model.sortformer_modules.fifo_len = 40
+        self.model.sortformer_modules.spkcache_update_period = 300
+        self.model.sortformer_modules.spkcache_len = 188
+        self.model.sortformer_modules._check_streaming_parameters()
+
+        warn(
+            "Sortformer supports maximum of 4 speakers only, "
+            "please use MSDD if your audio has more than 4 speakers",
+            Warning,
+        )
+
+        self.model.eval()
+
+    def diarize(self, audio: torch.Tensor):
+        with torch.inference_mode():
+            processed_signal, processed_signal_length = self.model.process_signal(
+                audio_signal=audio,
+                audio_signal_length=torch.tensor([audio.shape[-1]]),
+            )
+            processed_signal = processed_signal[:, :, : processed_signal_length.max()]
+
+            preds = self.model.forward_streaming(processed_signal, processed_signal_length)
+            preds = preds.cpu()
+
+        # TODO: make this tunable
+        diarize_cfg = DiarizeConfig(
+            postprocessing_params={
+                "onset": 0.5,
+                "offset": 0.5,
+                "pad_onset": 0.0,
+                "pad_offset": 0.0,
+                "min_duration_on": 0.0,
+                "min_duration_off": 0.0,
+            }
+        )
+
+        audio_rttm_map_dict = {
+            "audio": {
+                "uniq_id": "audio",
+                "audio_filepath": "tensor_audio",
+                "offset": 0.0,
+                "duration": None,
+                "text": "-",
+                "label": "infer",
+            }
+        }
+
+        self.model._diarize_audio_rttm_map = audio_rttm_map_dict
+        uniq_ids = list(self.model._diarize_audio_rttm_map.keys())
+        processed_outputs = self.model._diarize_output_processing(preds, uniq_ids, diarize_cfg)
+        self.model._diarize_audio_rttm_map = {}
+
+        labels = []
+        for label in processed_outputs[0]:
+            start, end, speaker = label.split()
+            start, end = float(start), float(end)
+            start, end = int(start * 1000), int(end * 1000)
+            labels.append((start, end, int(speaker.split("_")[1])))
+
+        labels = sorted(labels, key=lambda x: x[0])
+
+        return labels

diarize.py

@@ -89,7 +89,7 @@
 parser.add_argument(
     "--diarizer",
     default="msdd",
-    choices=["msdd"],
+    choices=["msdd", "sortformer"],
     help="Choose the diarization model to use",
 )
 
@@ -189,6 +189,11 @@
 
     diarizer_model = MSDDDiarizer(device=args.device)
 
+elif args.diarizer == "sortformer":
+    from diarization import SortformerDiarizer
+
+    diarizer_model = SortformerDiarizer(device=args.device)
+
 speaker_ts = diarizer_model.diarize(torch.from_numpy(audio_waveform).unsqueeze(0))
 del diarizer_model
 torch.cuda.empty_cache()

requirements.txt

@@ -1,4 +1,4 @@
-nemo_toolkit[asr]>=2.3.0
+nemo_toolkit[asr]>=2.5.0
 nltk
 faster-whisper>=1.1.0
 git+https://github.com/MahmoudAshraf97/demucs.git