Add real-time speaker diarization support

- New whisper_live/diarization.py: SpeakerDiarizer with online clustering
- Uses pyannote.audio speaker embeddings (optional dependency)
- Cosine similarity threshold for speaker matching (default 0.55)
- Running average embedding update for speaker stability
- Configurable max_speakers limit (default 10)
- Client options: enable_diarization, max_speakers
- Segments include 'speaker' field when diarization is active
- Graceful fallback: logs warning if pyannote not installed
- Added 12 unit tests (mock-based, no GPU required)

tests/test_client.py

@@ -56,6 +56,8 @@ class TestClientCallbacks(BaseTestCase):
             "enable_translation": False,
             "target_language": "fr",
             "hotwords": None,
+            "enable_diarization": False,
+            "max_speakers": 10,
         })
         self.client.on_open(self.mock_ws_app)
         self.mock_ws_app.send.assert_called_with(expected_message)

tests/test_diarization.py

@@ -0,0 +1,152 @@
+import unittest
+from unittest.mock import MagicMock, patch
+import numpy as np
+
+
+class TestSpeakerDiarizer(unittest.TestCase):
+    """Tests for SpeakerDiarizer with mocked embedding model."""
+
+    def _make_diarizer(self, **kwargs):
+        from whisper_live.diarization import SpeakerDiarizer
+        d = SpeakerDiarizer(**kwargs)
+        # Mock the embedding model to return deterministic embeddings
+        d._model = MagicMock()
+        return d
+
+    def _set_embedding(self, diarizer, embedding):
+        """Configure mock model to return a specific embedding."""
+        emb = np.array(embedding, dtype=np.float32)
+        emb = emb / np.linalg.norm(emb)
+        diarizer._model.return_value = emb
+
+    def test_first_speaker_creates_new(self):
+        d = self._make_diarizer()
+        self._set_embedding(d, [1.0, 0.0, 0.0])
+        audio = np.zeros(16000, dtype=np.float32)  # 1 second of audio
+        speaker = d.identify_speaker(audio)
+        self.assertEqual(speaker, "SPEAKER_00")
+        self.assertEqual(len(d.speakers), 1)
+
+    def test_same_speaker_matches(self):
+        d = self._make_diarizer(similarity_threshold=0.8)
+        self._set_embedding(d, [1.0, 0.0, 0.0])
+        audio = np.zeros(16000, dtype=np.float32)
+        d.identify_speaker(audio)  # SPEAKER_00
+        # Same embedding should match
+        self._set_embedding(d, [0.99, 0.01, 0.0])
+        speaker = d.identify_speaker(audio)
+        self.assertEqual(speaker, "SPEAKER_00")
+        self.assertEqual(len(d.speakers), 1)
+
+    def test_different_speaker_creates_new(self):
+        d = self._make_diarizer(similarity_threshold=0.8)
+        self._set_embedding(d, [1.0, 0.0, 0.0])
+        audio = np.zeros(16000, dtype=np.float32)
+        d.identify_speaker(audio)  # SPEAKER_00
+
+        # Very different embedding
+        self._set_embedding(d, [0.0, 1.0, 0.0])
+        speaker = d.identify_speaker(audio)
+        self.assertEqual(speaker, "SPEAKER_01")
+        self.assertEqual(len(d.speakers), 2)
+
+    def test_max_speakers_limit(self):
+        d = self._make_diarizer(similarity_threshold=0.95, max_speakers=2)
+        audio = np.zeros(16000, dtype=np.float32)
+
+        self._set_embedding(d, [1.0, 0.0, 0.0])
+        d.identify_speaker(audio)  # SPEAKER_00
+        self._set_embedding(d, [0.0, 1.0, 0.0])
+        d.identify_speaker(audio)  # SPEAKER_01
+
+        # Third distinct speaker should be assigned to closest existing
+        self._set_embedding(d, [0.0, 0.0, 1.0])
+        speaker = d.identify_speaker(audio)
+        self.assertIn(speaker, ["SPEAKER_00", "SPEAKER_01"])
+        self.assertEqual(len(d.speakers), 2)
+
+    def test_short_audio_returns_none(self):
+        d = self._make_diarizer()
+        # Less than 0.3 seconds
+        audio = np.zeros(3000, dtype=np.float32)
+        speaker = d.identify_speaker(audio)
+        self.assertIsNone(speaker)
+
+    def test_reset_clears_state(self):
+        d = self._make_diarizer()
+        self._set_embedding(d, [1.0, 0.0, 0.0])
+        audio = np.zeros(16000, dtype=np.float32)
+        d.identify_speaker(audio)
+        self.assertEqual(len(d.speakers), 1)
+        d.reset()
+        self.assertEqual(len(d.speakers), 0)
+        self.assertEqual(d._speaker_count, 0)
+
+    def test_import_error_without_pyannote(self):
+        from whisper_live.diarization import SpeakerDiarizer
+        d = SpeakerDiarizer()
+        with patch.dict("sys.modules", {"pyannote": None, "pyannote.audio": None}):
+            with self.assertRaises(ImportError):
+                d._load_model()
+
+
+class TestDiarizationInBase(unittest.TestCase):
+    """Test diarization integration in ServeClientBase."""
+
+    def _make_client(self, diarization=None):
+        from whisper_live.backend.base import ServeClientBase
+
+        class ConcreteClient(ServeClientBase):
+            def __init__(self, **kwargs):
+                super().__init__(**kwargs)
+                self.language = "en"
+            def transcribe_audio(self, input_sample):
+                return None
+            def handle_transcription_output(self, result, duration):
+                pass
+
+        ws = MagicMock()
+        return ConcreteClient(
+            client_uid="test-uid", websocket=ws, diarization=diarization
+        )
+
+    def test_no_diarization_by_default(self):
+        client = self._make_client()
+        self.assertIsNone(client.diarization)
+
+    def test_format_segment_with_speaker(self):
+        client = self._make_client()
+        seg = client.format_segment(0.0, 1.0, "hello", speaker="SPEAKER_00")
+        self.assertEqual(seg["speaker"], "SPEAKER_00")
+
+    def test_format_segment_without_speaker(self):
+        client = self._make_client()
+        seg = client.format_segment(0.0, 1.0, "hello")
+        self.assertNotIn("speaker", seg)
+
+    def test_identify_speaker_disabled(self):
+        client = self._make_client(diarization=None)
+        seg = MagicMock()
+        seg.start = 0.0
+        seg.end = 1.0
+        result = client._identify_speaker(seg)
+        self.assertIsNone(result)
+
+    def test_identify_speaker_calls_diarizer(self):
+        mock_diarizer = MagicMock()
+        mock_diarizer.identify_speaker.return_value = "SPEAKER_01"
+        client = self._make_client(diarization=mock_diarizer)
+        # Set up audio buffer
+        client.frames_np = np.zeros(48000, dtype=np.float32)
+        client.frames_offset = 0.0
+        client.timestamp_offset = 0.0
+        seg = MagicMock()
+        seg.start = 0.5
+        seg.end = 1.5
+        result = client._identify_speaker(seg)
+        self.assertEqual(result, "SPEAKER_01")
+        mock_diarizer.identify_speaker.assert_called_once()
+
+
+if __name__ == "__main__":
+    unittest.main()

whisper_live/backend/base.py

@@ -38,6 +38,7 @@ class ServeClientBase(object):
         clip_audio=False,
         same_output_threshold=10,
         translation_queue=None,
+        diarization=None,
     ):
         self.client_uid = client_uid
         self.websocket = websocket
@@ -45,6 +46,7 @@ class ServeClientBase(object):
         self.no_speech_thresh = no_speech_thresh
         self.clip_audio = clip_audio
         self.same_output_threshold = same_output_threshold
+        self.diarization = diarization
 
         self.frames = b""
         self.timestamp_offset = 0.0
@@ -116,7 +118,7 @@ class ServeClientBase(object):
     def handle_transcription_output(self, result, duration):
         raise NotImplementedError
     
-    def format_segment(self, start, end, text, completed=False):
+    def format_segment(self, start, end, text, completed=False, speaker=None):
         """
         Formats a transcription segment with precise start and end times alongside the transcribed text.
 
@@ -124,18 +126,22 @@ class ServeClientBase(object):
             start (float): The start time of the transcription segment in seconds.
             end (float): The end time of the transcription segment in seconds.
             text (str): The transcribed text corresponding to the segment.
+            speaker (str, optional): Speaker label from diarization.
 
         Returns:
             dict: A dictionary representing the formatted transcription segment, including
                 'start' and 'end' times as strings with three decimal places and the 'text'
                 of the transcription.
         """
-        return {
+        seg = {
             'start': "{:.3f}".format(start),
             'end': "{:.3f}".format(end),
             'text': text,
-            'completed': completed
+            'completed': completed,
         }
+        if speaker is not None:
+            seg['speaker'] = speaker
+        return seg
 
     def add_frames(self, frame_np):
         """
@@ -292,6 +298,29 @@ class ServeClientBase(object):
     def get_segment_end(self, segment):
         return getattr(segment, "end", getattr(segment, "end_ts", 0))
 
+    def _identify_speaker(self, segment):
+        """Run diarization on a segment's audio slice if diarization is enabled.
+
+        Returns:
+            str or None: Speaker label, or None if diarization is disabled or audio unavailable.
+        """
+        if self.diarization is None or self.frames_np is None:
+            return None
+        try:
+            seg_start = self.get_segment_start(segment)
+            seg_end = self.get_segment_end(segment)
+            start_sample = int(seg_start * self.RATE)
+            end_sample = int(seg_end * self.RATE)
+            # Extract audio relative to the current buffer
+            samples_offset = max(0, int((self.timestamp_offset - self.frames_offset) * self.RATE))
+            audio_slice = self.frames_np[samples_offset + start_sample:samples_offset + end_sample]
+            if len(audio_slice) < self.RATE * 0.3:
+                return None
+            return self.diarization.identify_speaker(audio_slice, self.RATE)
+        except Exception as e:
+            logging.error(f"Diarization error: {e}")
+            return None
+
     def update_segments(self, segments, duration):
         """
         Processes the segments from Whisper and updates the transcript.
@@ -321,7 +350,8 @@ class ServeClientBase(object):
                     continue
                 if self.get_segment_no_speech_prob(s) > self.no_speech_thresh:
                     continue
-                completed_segment = self.format_segment(start, end, text_, completed=True)
+                speaker = self._identify_speaker(s)
+                completed_segment = self.format_segment(start, end, text_, completed=True, speaker=speaker)
                 self.transcript.append(completed_segment)
 
                 if self.translation_queue:

whisper_live/backend/faster_whisper_backend.py

@@ -35,6 +35,7 @@ class ServeClientFasterWhisper(ServeClientBase):
         cache_path="~/.cache/whisper-live/",
         translation_queue=None,
         hotwords=None,
+        diarization=None,
     ):
         """
         Initialize a ServeClient instance.
@@ -64,7 +65,8 @@ class ServeClientFasterWhisper(ServeClientBase):
             no_speech_thresh,
             clip_audio,
             same_output_threshold,
-            translation_queue
+            translation_queue,
+            diarization,
         )
         self.cache_path = cache_path
         self.model_sizes = [

whisper_live/client.py

@@ -44,6 +44,8 @@ class Client:
         enable_timestamps=False,
         display_segments=4,
         hotwords=None,
+        enable_diarization=False,
+        max_speakers=10,
     ):
         """
         Initializes a Client instance for audio recording and streaming to a server.
@@ -103,7 +105,8 @@ class Client:
         self.enable_timestamps = enable_timestamps
         self.display_segments = display_segments
         self.hotwords = hotwords
-
+        self.enable_diarization = enable_diarization
+        self.max_speakers = max_speakers
         self.audio_bytes = None
 
         if host is not None and port is not None:
@@ -302,6 +305,8 @@ class Client:
                     "enable_translation": self.enable_translation,
                     "target_language": self.target_language,
                     "hotwords": self.hotwords,
+                    "enable_diarization": self.enable_diarization,
+                    "max_speakers": self.max_speakers,
                 }
             )
         )
@@ -824,7 +829,10 @@ class TranscriptionClient(TranscriptionTeeClient):
         enable_timestamps=False,
         display_segments=4,
         hotwords=None,
+        enable_diarization=False,
+        max_speakers=10,
     ):
+
         self.client = Client(
             host,
             port,
@@ -847,6 +855,8 @@ class TranscriptionClient(TranscriptionTeeClient):
             enable_timestamps=enable_timestamps,
             display_segments=display_segments,
             hotwords=hotwords,
+            enable_diarization=enable_diarization,
+            max_speakers=max_speakers,
         )
 
         if save_output_recording and not output_recording_filename.endswith(".wav"):

whisper_live/diarization.py

@@ -0,0 +1,142 @@
+"""
+Optional speaker diarization module for WhisperLive.
+
+Uses speaker embeddings and online clustering to assign speaker labels
+to transcription segments in real-time. Requires pyannote.audio as an
+optional dependency.
+
+Install: pip install pyannote.audio
+"""
+
+import logging
+import numpy as np
+
+
+class SpeakerDiarizer:
+    """Real-time speaker diarization using speaker embeddings and online clustering.
+
+    Each completed transcription segment's audio is passed through a speaker
+    embedding model. The embedding is compared against known speakers using
+    cosine similarity. If no match exceeds the threshold, a new speaker is
+    created.
+
+    Args:
+        similarity_threshold (float): Minimum cosine similarity to match an
+            existing speaker. Lower values merge speakers more aggressively.
+            Default 0.55.
+        max_speakers (int): Maximum number of distinct speakers to track.
+            Once reached, new segments are assigned to the closest existing
+            speaker. Default 10.
+        embedding_model (str): The pyannote embedding model to use.
+            Default "pyannote/wespeaker-voxceleb-resnet34-LM".
+        hf_token (str or None): HuggingFace token for gated model access.
+    """
+
+    def __init__(
+        self,
+        similarity_threshold=0.55,
+        max_speakers=10,
+        embedding_model="pyannote/wespeaker-voxceleb-resnet34-LM",
+        hf_token=None,
+    ):
+        self.similarity_threshold = similarity_threshold
+        self.max_speakers = max_speakers
+        self.speakers = {}  # speaker_id -> embedding (averaged)
+        self._speaker_count = 0
+        self._model = None
+        self._embedding_model_name = embedding_model
+        self._hf_token = hf_token
+
+    def _load_model(self):
+        """Lazy-load the embedding model on first use."""
+        if self._model is not None:
+            return
+        try:
+            from pyannote.audio import Model, Inference
+            import torch
+
+            model = Model.from_pretrained(
+                self._embedding_model_name,
+                use_auth_token=self._hf_token,
+            )
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            self._model = Inference(model, window="whole", device=torch.device(device))
+            logging.info(f"Speaker embedding model loaded on {device}")
+        except ImportError:
+            raise ImportError(
+                "pyannote.audio is required for speaker diarization. "
+                "Install it with: pip install pyannote.audio"
+            )
+
+    def _compute_embedding(self, audio_np, sample_rate=16000):
+        """Compute a speaker embedding from an audio numpy array.
+
+        Args:
+            audio_np (np.ndarray): 1-D float32 audio samples.
+            sample_rate (int): Sample rate of the audio.
+
+        Returns:
+            np.ndarray: Speaker embedding vector, or None if audio is too short.
+        """
+        self._load_model()
+        if len(audio_np) < sample_rate * 0.3:
+            return None
+        waveform = {
+            "waveform": __import__("torch").tensor(audio_np).unsqueeze(0),
+            "sample_rate": sample_rate,
+        }
+        embedding = self._model(waveform)
+        return embedding / np.linalg.norm(embedding)
+
+    @staticmethod
+    def _cosine_similarity(a, b):
+        """Compute cosine similarity between two vectors."""
+        return float(np.dot(a, b))
+
+    def identify_speaker(self, audio_np, sample_rate=16000):
+        """Identify or create a speaker from an audio segment.
+
+        Args:
+            audio_np (np.ndarray): 1-D float32 audio for the segment.
+            sample_rate (int): Sample rate. Default 16000.
+
+        Returns:
+            str or None: Speaker label (e.g. "SPEAKER_00"), or None if
+                the audio is too short to embed.
+        """
+        embedding = self._compute_embedding(audio_np, sample_rate)
+        if embedding is None:
+            return None
+
+        best_speaker = None
+        best_sim = -1.0
+
+        for speaker_id, stored_emb in self.speakers.items():
+            sim = self._cosine_similarity(embedding, stored_emb)
+            if sim > best_sim:
+                best_sim = sim
+                best_speaker = speaker_id
+
+        if best_sim >= self.similarity_threshold:
+            # Update running average for the matched speaker
+            self.speakers[best_speaker] = (
+                self.speakers[best_speaker] * 0.9 + embedding * 0.1
+            )
+            # Re-normalize
+            self.speakers[best_speaker] /= np.linalg.norm(self.speakers[best_speaker])
+            return best_speaker
+
+        if len(self.speakers) >= self.max_speakers:
+            # Assign to closest speaker
+            return best_speaker if best_speaker else f"SPEAKER_{self._speaker_count:02d}"
+
+        # Create a new speaker
+        speaker_id = f"SPEAKER_{self._speaker_count:02d}"
+        self._speaker_count += 1
+        self.speakers[speaker_id] = embedding
+        return speaker_id
+
+    def reset(self):
+        """Reset all speaker state."""
+        self.speakers.clear()
+        self._speaker_count = 0

whisper_live/server.py

@@ -292,6 +292,7 @@ class TranscriptionServer:
                     cache_path=self.cache_path,
                     translation_queue=translation_queue,
                     hotwords=options.get("hotwords"),
+                    diarization=self._create_diarizer(options),
                 )
 
                 logging.info("Running faster_whisper backend.")
@@ -320,6 +321,25 @@ class TranscriptionServer:
 
         self.client_manager.add_client(websocket, client)
 
+    def _create_diarizer(self, options):
+        """Create a SpeakerDiarizer if the client requested diarization.
+
+        Returns:
+            SpeakerDiarizer or None
+        """
+        if not options.get("enable_diarization", False):
+            return None
+        try:
+            from whisper_live.diarization import SpeakerDiarizer
+            return SpeakerDiarizer(
+                similarity_threshold=options.get("diarization_threshold", 0.55),
+                max_speakers=options.get("max_speakers", 10),
+                hf_token=options.get("hf_token"),
+            )
+        except ImportError:
+            logging.warning("pyannote.audio not installed; diarization disabled")
+            return None
+
     def get_audio_from_websocket(self, websocket):
         """
         Receives audio buffer from websocket and creates a numpy array out of it.