Support known speaker hints in REST API

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()

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,