cjones/WhisperLive
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add real-time speaker diarization support

Browse Source 
- 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)
This commit is contained in:
Aaron Boxer
2026-04-17 10:25:10 -04:00
committed by Aaron Boxer
parent 3d63e82571
commit 18b897277f
7 changed files with 364 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
whisper_live/backend/base.py
+34 -4
View File
@@ -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
+3 -1
View File
@@ -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
+11 -1
View File
@@ -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
+142
View File
@@ -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
+20
View File
@@ -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.