diff --git a/README.md b/README.md index 9ece51f..09416da 100644 --- a/README.md +++ b/README.md @@ -245,6 +245,8 @@ When enabled, completed segments include a `speaker` field: ``` Diarization uses online cosine-similarity clustering of speaker embeddings. If `pyannote.audio` is not installed, the server logs a warning and continues without diarization. +The OpenAI-compatible REST endpoint also accepts `known_speaker_names` and uploaded `known_speaker_references` multipart fields. When speaker fields are supplied with `response_format="verbose_json"`, segments include a `speaker` field. + #### Batch Inference Batch multiple client sessions into single GPU calls for higher throughput: ```bash diff --git a/tests/test_diarization.py b/tests/test_diarization.py index 00957b1..dd675c3 100644 --- a/tests/test_diarization.py +++ b/tests/test_diarization.py @@ -8,6 +8,7 @@ class TestSpeakerDiarizer(unittest.TestCase): 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() @@ -82,8 +83,26 @@ class TestSpeakerDiarizer(unittest.TestCase): self.assertEqual(len(d.speakers), 0) self.assertEqual(d._speaker_count, 0) + def test_enroll_speaker_uses_known_name(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) + self.assertTrue(d.enroll_speaker("Alice", audio)) + + self._set_embedding(d, [0.99, 0.01, 0.0]) + speaker = d.identify_speaker(audio) + self.assertEqual(speaker, "Alice") + + def test_speaker_names_label_new_speakers(self): + d = self._make_diarizer(speaker_names=["Alice"]) + self._set_embedding(d, [1.0, 0.0, 0.0]) + audio = np.zeros(16000, dtype=np.float32) + speaker = d.identify_speaker(audio) + self.assertEqual(speaker, "Alice") + 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): @@ -100,8 +119,10 @@ class TestDiarizationInBase(unittest.TestCase): 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 @@ -148,5 +169,33 @@ class TestDiarizationInBase(unittest.TestCase): mock_diarizer.identify_speaker.assert_called_once() +class TestRestDiarizationHelpers(unittest.TestCase): + def test_normalize_form_list_accepts_repeated_or_comma_separated_values(self): + from whisper_live.server import TranscriptionServer + + self.assertEqual( + TranscriptionServer._normalize_form_list(["Alice,Bob", "Carol"]), + ["Alice", "Bob", "Carol"], + ) + + def test_speaker_labels_for_segments(self): + from whisper_live.server import TranscriptionServer + + segment = MagicMock() + segment.start = 0.0 + segment.end = 1.0 + diarizer = MagicMock() + diarizer.identify_speaker.return_value = "Alice" + + speakers = TranscriptionServer._speaker_labels_for_segments( + [segment], + np.zeros(16000, dtype=np.float32), + diarizer, + ) + + self.assertEqual(speakers, {0: "Alice"}) + diarizer.identify_speaker.assert_called_once() + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_server_extended.py b/tests/test_server_extended.py index 419755e..3a59eb9 100644 --- a/tests/test_server_extended.py +++ b/tests/test_server_extended.py @@ -343,11 +343,6 @@ class TestStreamTranscription(unittest.TestCase): def _make_app(self): """Create a FastAPI app with the transcribe endpoint that has streaming support.""" from fastapi import FastAPI, UploadFile, Form - from fastapi.testclient import TestClient - from starlette.responses import StreamingResponse - import os - import tempfile - import shutil app = FastAPI() server = TranscriptionServer() @@ -504,10 +499,9 @@ class TestRESTAPIParamWarnings(unittest.TestCase): def setUpClass(cls): """Build a FastAPI test app by extracting the endpoint definition.""" import logging - from fastapi import FastAPI, UploadFile, Form + from fastapi import FastAPI, UploadFile, Form, File from fastapi.testclient import TestClient from typing import Optional, List - from starlette.responses import PlainTextResponse, JSONResponse app = FastAPI() @@ -523,16 +517,12 @@ class TestRESTAPIParamWarnings(unittest.TestCase): chunking_strategy: Optional[str] = Form(default=None), include: Optional[List[str]] = Form(default=None), known_speaker_names: Optional[List[str]] = Form(default=None), - known_speaker_references: Optional[List[str]] = Form(default=None), + known_speaker_references: Optional[List[UploadFile]] = File(default=None), stream: bool = Form(default=False), ): ignored_params = [] if chunking_strategy: ignored_params.append(f"chunking_strategy='{chunking_strategy}'") - if known_speaker_names: - ignored_params.append("known_speaker_names") - if known_speaker_references: - ignored_params.append("known_speaker_references") if include: ignored_params.append(f"include={include}") if ignored_params: @@ -565,17 +555,17 @@ class TestRESTAPIParamWarnings(unittest.TestCase): ignored = resp.json()["ignored"] self.assertTrue(any("include" in p for p in ignored)) - def test_known_speaker_names_warning(self): + def test_known_speaker_names_supported(self): resp = self._post(known_speaker_names="alice") self.assertEqual(resp.status_code, 200) ignored = resp.json()["ignored"] - self.assertTrue(any("known_speaker_names" in p for p in ignored)) + self.assertFalse(any("known_speaker_names" in p for p in ignored)) def test_multiple_ignored_params(self): resp = self._post(chunking_strategy="auto", known_speaker_names="bob") self.assertEqual(resp.status_code, 200) ignored = resp.json()["ignored"] - self.assertGreaterEqual(len(ignored), 2) + self.assertEqual(len(ignored), 1) class TestAPIKeyAuth(unittest.TestCase): diff --git a/whisper_live/diarization.py b/whisper_live/diarization.py index c58bf58..4141aca 100644 --- a/whisper_live/diarization.py +++ b/whisper_live/diarization.py @@ -12,6 +12,28 @@ import logging import numpy as np +def load_audio(file_path, sample_rate=16000): + """Load an audio file as mono float32 PCM at the requested sample rate.""" + import av + + container = av.open(file_path) + resampler = av.AudioResampler(format="flt", layout="mono", rate=sample_rate) + chunks = [] + + try: + for frame in container.decode(audio=0): + for resampled_frame in resampler.resample(frame): + chunks.append( + resampled_frame.to_ndarray().reshape(-1).astype(np.float32) + ) + finally: + container.close() + + if not chunks: + return np.array([], dtype=np.float32) + return np.concatenate(chunks) + + class SpeakerDiarizer: """Real-time speaker diarization using speaker embeddings and online clustering. @@ -38,15 +60,22 @@ class SpeakerDiarizer: max_speakers=10, embedding_model="pyannote/wespeaker-voxceleb-resnet34-LM", hf_token=None, + speaker_names=None, ): self.similarity_threshold = similarity_threshold self.max_speakers = max_speakers + self.speaker_names = list(speaker_names or []) 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 _next_speaker_id(self): + if self._speaker_count < len(self.speaker_names): + return self.speaker_names[self._speaker_count] + return f"SPEAKER_{self._speaker_count:02d}" + def _load_model(self): """Lazy-load the embedding model on first use.""" if self._model is not None: @@ -128,14 +157,24 @@ class SpeakerDiarizer: if len(self.speakers) >= self.max_speakers: # Assign to closest speaker - return best_speaker if best_speaker else f"SPEAKER_{self._speaker_count:02d}" + 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}" + speaker_id = self._next_speaker_id() self._speaker_count += 1 self.speakers[speaker_id] = embedding return speaker_id + def enroll_speaker(self, speaker_name, audio_np, sample_rate=16000): + """Enroll a known speaker from reference audio.""" + embedding = self._compute_embedding(audio_np, sample_rate) + if embedding is None: + return False + self.speakers[speaker_name] = embedding + return True + def reset(self): """Reset all speaker state.""" self.speakers.clear() diff --git a/whisper_live/server.py b/whisper_live/server.py index 8504683..d6c7ebc 100644 --- a/whisper_live/server.py +++ b/whisper_live/server.py @@ -9,19 +9,18 @@ import logging import shutil import tempfile from typing import Optional, List -from fastapi import FastAPI, UploadFile, Form, Request +from fastapi import FastAPI, UploadFile, Form, Request, File from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse -from starlette.responses import PlainTextResponse, JSONResponse, StreamingResponse +from starlette.responses import PlainTextResponse, StreamingResponse import uvicorn from faster_whisper import WhisperModel import torch from enum import Enum -from whisper_live import metrics as wl_metrics -from typing import List, Optional import numpy as np +from whisper_live import metrics as wl_metrics from websockets.sync.server import serve from websockets.exceptions import ConnectionClosed from whisper_live.vad import VoiceActivityDetector @@ -523,6 +522,67 @@ class TranscriptionServer: return StreamingResponse(_sse_generator(), media_type="text/event-stream") + @staticmethod + def _normalize_form_list(values): + """Normalize repeated or comma-separated multipart form fields.""" + if not values: + return [] + normalized = [] + for value in values: + if isinstance(value, str): + normalized.extend(item.strip() for item in value.split(",") if item.strip()) + return normalized + + async def _create_rest_diarizer(self, known_speaker_names, known_speaker_references): + """Create a diarizer from OpenAI-compatible known speaker fields.""" + speaker_names = self._normalize_form_list(known_speaker_names) + speaker_references = known_speaker_references or [] + + if speaker_references and not speaker_names: + raise ValueError("known_speaker_references requires matching known_speaker_names") + if speaker_names and speaker_references and len(speaker_names) != len(speaker_references): + raise ValueError("known_speaker_names and known_speaker_references must have the same length") + if not speaker_names and not speaker_references: + return None + + from whisper_live.diarization import SpeakerDiarizer, load_audio + + diarizer = SpeakerDiarizer( + max_speakers=max(10, len(speaker_names)), + speaker_names=speaker_names, + ) + + for speaker_name, reference in zip(speaker_names, speaker_references): + suffix = os.path.splitext(reference.filename or "")[1] or ".wav" + reference_path = None + try: + with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: + tmp.write(await reference.read()) + reference_path = tmp.name + audio_np = load_audio(reference_path) + if not diarizer.enroll_speaker(speaker_name, audio_np): + raise ValueError(f"known_speaker_references for '{speaker_name}' is too short") + finally: + if reference_path and os.path.exists(reference_path): + os.unlink(reference_path) + + return diarizer + + @staticmethod + def _speaker_labels_for_segments(segments, audio_np, diarizer, sample_rate=16000): + if diarizer is None or audio_np is None: + return {} + labels = {} + for index, segment in enumerate(segments): + start = max(0, int(segment.start * sample_rate)) + end = min(len(audio_np), int(segment.end * sample_rate)) + if end <= start: + continue + speaker = diarizer.identify_speaker(audio_np[start:end], sample_rate) + if speaker: + labels[index] = speaker + return labels + def run(self, host, port=9090, @@ -666,7 +726,7 @@ class TranscriptionServer: chunking_strategy: Optional[str] = Form(default=None), include: Optional[List[str]] = Form(default=None), known_speaker_names: Optional[List[str]] = Form(default=None), - known_speaker_references: Optional[List[str]] = Form(default=None), + known_speaker_references: Optional[List[UploadFile]] = File(default=None), stream: bool = Form(default=False), hotwords: Optional[str] = Form(default=None), ): @@ -680,10 +740,6 @@ class TranscriptionServer: ignored_params = [] if chunking_strategy: ignored_params.append(f"chunking_strategy='{chunking_strategy}'") - if known_speaker_names: - ignored_params.append("known_speaker_names") - if known_speaker_references: - ignored_params.append("known_speaker_references") if include: ignored_params.append(f"include={include}") if ignored_params: @@ -698,6 +754,7 @@ class TranscriptionServer: logging.warning(f"Model '{model}' requested; using 'small' as fallback.") model_name = faster_whisper_custom_model_path or "small" + tmp_path = None try: suffix = os.path.splitext(file.filename)[1] or ".wav" with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: @@ -717,9 +774,9 @@ class TranscriptionServer: word_timestamps=(timestamp_granularities and "word" in timestamp_granularities), hotwords=hotwords, ) + segments = list(segments) text = " ".join([s.text.strip() for s in segments]) - os.unlink(tmp_path) if response_format == "text": wl_metrics.track_rest_request(endpoint="transcriptions", status=200) @@ -735,7 +792,17 @@ class TranscriptionServer: "text": text, "segments": [] } - for seg in segments: + speaker_labels = {} + try: + rest_diarizer = await self._create_rest_diarizer(known_speaker_names, known_speaker_references) + except ValueError as e: + wl_metrics.track_rest_request(endpoint="transcriptions", status=400) + return JSONResponse({"error": str(e)}, status_code=400) + if rest_diarizer is not None: + from whisper_live.diarization import load_audio + audio_np = load_audio(tmp_path) + speaker_labels = self._speaker_labels_for_segments(segments, audio_np, rest_diarizer) + for index, seg in enumerate(segments): seg_dict = { "id": seg.id, "seek": seg.seek, @@ -748,6 +815,8 @@ class TranscriptionServer: "compression_ratio": seg.compression_ratio, "no_speech_prob": seg.no_speech_prob } + if index in speaker_labels: + seg_dict["speaker"] = speaker_labels[index] if timestamp_granularities and "word" in timestamp_granularities: seg_dict["words"] = [{"word": w.word, "start": w.start, "end": w.end, "probability": w.probability} for w in seg.words] verbose["segments"].append(seg_dict) @@ -768,6 +837,9 @@ class TranscriptionServer: wl_metrics.track_rest_request(endpoint="transcriptions", status=500) wl_metrics.track_error("rest_transcription") return JSONResponse({"error": str(e)}, status_code=500) + finally: + if tmp_path and os.path.exists(tmp_path): + os.unlink(tmp_path) threading.Thread( target=uvicorn.run,