From 18b897277f68b943fc63acc8450903094ffbea91 Mon Sep 17 00:00:00 2001 From: Aaron Boxer Date: Fri, 17 Apr 2026 10:25:10 -0400 Subject: [PATCH] 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 | 2 + tests/test_diarization.py | 152 ++++++++++++++++++ whisper_live/backend/base.py | 38 ++++- .../backend/faster_whisper_backend.py | 4 +- whisper_live/client.py | 12 +- whisper_live/diarization.py | 142 ++++++++++++++++ whisper_live/server.py | 20 +++ 7 files changed, 364 insertions(+), 6 deletions(-) create mode 100644 tests/test_diarization.py create mode 100644 whisper_live/diarization.py diff --git a/tests/test_client.py b/tests/test_client.py index b636000..d55bf6d 100644 --- a/tests/test_client.py +++ b/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) diff --git a/tests/test_diarization.py b/tests/test_diarization.py new file mode 100644 index 0000000..00957b1 --- /dev/null +++ b/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() diff --git a/whisper_live/backend/base.py b/whisper_live/backend/base.py index 78182c8..1926730 100644 --- a/whisper_live/backend/base.py +++ b/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: diff --git a/whisper_live/backend/faster_whisper_backend.py b/whisper_live/backend/faster_whisper_backend.py index 76e498b..02b07e4 100644 --- a/whisper_live/backend/faster_whisper_backend.py +++ b/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 = [ diff --git a/whisper_live/client.py b/whisper_live/client.py index 14af63f..c8ac4e5 100644 --- a/whisper_live/client.py +++ b/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"): diff --git a/whisper_live/diarization.py b/whisper_live/diarization.py new file mode 100644 index 0000000..c58bf58 --- /dev/null +++ b/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 diff --git a/whisper_live/server.py b/whisper_live/server.py index 32d24aa..66f9324 100644 --- a/whisper_live/server.py +++ b/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.