416 lines
15 KiB
Python
416 lines
15 KiB
Python
"""MeetVault local live speech engine.
|
|
|
|
Protocol
|
|
--------
|
|
WebSocket: ws://127.0.0.1:8765/ws/live
|
|
1. Client sends a JSON config message.
|
|
2. Client streams signed 16-bit little-endian mono PCM at 16 kHz.
|
|
3. Server returns JSON events: ready, partial, final, health, error.
|
|
|
|
Real mode combines faster-whisper with diart speaker diarization. Mock mode is
|
|
available so the desktop UI can be exercised without ML dependencies/models.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import queue
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
import uvicorn
|
|
|
|
SAMPLE_RATE = 16000
|
|
|
|
|
|
@dataclass
|
|
class EngineConfig:
|
|
model: str = "small"
|
|
compute: str = "auto"
|
|
language: str = "auto"
|
|
diarization: bool = True
|
|
sample_rate: int = SAMPLE_RATE
|
|
|
|
|
|
class SpeakerTimeline:
|
|
"""Thread-safe speaker timeline built from diart streaming predictions."""
|
|
|
|
def __init__(self) -> None:
|
|
self._segments: list[tuple[float, float, str]] = []
|
|
self._label_map: dict[str, str] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
def _stable_label(self, raw: str) -> str:
|
|
with self._lock:
|
|
if raw not in self._label_map:
|
|
self._label_map[raw] = f"speaker_{len(self._label_map) + 1}"
|
|
return self._label_map[raw]
|
|
|
|
def update_from_annotation(self, annotation) -> None:
|
|
fresh: list[tuple[float, float, str]] = []
|
|
try:
|
|
for segment, _, label in annotation.itertracks(yield_label=True):
|
|
fresh.append((float(segment.start), float(segment.end), self._stable_label(str(label))))
|
|
except Exception:
|
|
return
|
|
if not fresh:
|
|
return
|
|
latest_end = max(end for _, end, _ in fresh)
|
|
with self._lock:
|
|
# Keep a rolling two-minute history, then merge the latest prediction.
|
|
self._segments = [s for s in self._segments if s[1] >= latest_end - 120]
|
|
self._segments.extend(fresh)
|
|
|
|
def resolve(self, start: float, end: float) -> str:
|
|
with self._lock:
|
|
candidates = list(self._segments)
|
|
best_label = "speaker_1"
|
|
best_overlap = 0.0
|
|
for s, e, label in candidates:
|
|
overlap = max(0.0, min(end, e) - max(start, s))
|
|
if overlap > best_overlap:
|
|
best_overlap = overlap
|
|
best_label = label
|
|
return best_label
|
|
|
|
|
|
class QueueAudioSource:
|
|
"""diart AudioSource fed by PCM pushed from the WebSocket thread."""
|
|
|
|
def __init__(self, sample_rate: int):
|
|
from diart.sources import AudioSource
|
|
|
|
class _Source(AudioSource):
|
|
def __init__(self, sr: int):
|
|
super().__init__("meetvault-live", sr)
|
|
self.q: queue.Queue[Optional[np.ndarray]] = queue.Queue()
|
|
self.cursor = 0
|
|
|
|
@property
|
|
def is_regular(self):
|
|
return False
|
|
|
|
@property
|
|
def duration(self):
|
|
return None
|
|
|
|
def push(self, samples: np.ndarray):
|
|
self.q.put(samples.astype(np.float32, copy=False))
|
|
|
|
def close(self):
|
|
self.q.put(None)
|
|
|
|
def read(self):
|
|
from pyannote.core import SlidingWindow, SlidingWindowFeature
|
|
|
|
while True:
|
|
item = self.q.get()
|
|
if item is None:
|
|
self.stream.on_completed()
|
|
return
|
|
start = self.cursor / self.sample_rate
|
|
window = SlidingWindow(start=start, duration=1 / self.sample_rate, step=1 / self.sample_rate)
|
|
feature = SlidingWindowFeature(item.reshape(-1, 1), window)
|
|
self.cursor += len(item)
|
|
self.stream.on_next(feature)
|
|
|
|
self.source = _Source(sample_rate)
|
|
|
|
def push(self, samples: np.ndarray) -> None:
|
|
self.source.push(samples)
|
|
|
|
def close(self) -> None:
|
|
self.source.close()
|
|
|
|
|
|
class DiartWorker:
|
|
def __init__(self, enabled: bool, sample_rate: int = SAMPLE_RATE):
|
|
self.enabled = enabled
|
|
self.timeline = SpeakerTimeline()
|
|
self.audio_source: Optional[QueueAudioSource] = None
|
|
self.thread: Optional[threading.Thread] = None
|
|
self.error: Optional[str] = None
|
|
if enabled:
|
|
self._start(sample_rate)
|
|
|
|
def _start(self, sample_rate: int) -> None:
|
|
try:
|
|
from diart import SpeakerDiarization
|
|
from diart.inference import StreamingInference
|
|
|
|
pipeline = SpeakerDiarization()
|
|
expected_rate = int(getattr(pipeline.config, "sample_rate", sample_rate))
|
|
if expected_rate != sample_rate:
|
|
raise RuntimeError(f"diart expects {expected_rate} Hz but MeetVault sends {sample_rate} Hz")
|
|
self.audio_source = QueueAudioSource(sample_rate)
|
|
inference = StreamingInference(pipeline, self.audio_source.source, do_plot=False)
|
|
|
|
def hook(annotation_and_waveform):
|
|
try:
|
|
self.timeline.update_from_annotation(annotation_and_waveform[0])
|
|
except Exception:
|
|
pass
|
|
|
|
inference.attach_hooks(hook)
|
|
self.thread = threading.Thread(target=inference, name="meetvault-diart", daemon=True)
|
|
self.thread.start()
|
|
except Exception as exc:
|
|
self.error = str(exc)
|
|
self.enabled = False
|
|
|
|
def push(self, samples: np.ndarray) -> None:
|
|
if self.enabled and self.audio_source:
|
|
self.audio_source.push(samples)
|
|
|
|
def close(self) -> None:
|
|
if self.audio_source:
|
|
self.audio_source.close()
|
|
|
|
def resolve(self, start: float, end: float) -> str:
|
|
return self.timeline.resolve(start, end) if self.enabled else "speaker_1"
|
|
|
|
|
|
class LiveSpeechSession:
|
|
def __init__(self, config: EngineConfig):
|
|
from faster_whisper import WhisperModel
|
|
|
|
self.config = config
|
|
device = config.compute
|
|
if device == "auto":
|
|
device = "cpu"
|
|
compute_type = "float16" if device == "cuda" else "int8"
|
|
try:
|
|
self.model = WhisperModel(config.model, device=device, compute_type=compute_type)
|
|
except Exception:
|
|
# Useful fallback when CUDA was selected but is not available.
|
|
self.model = WhisperModel(config.model, device="cpu", compute_type="int8")
|
|
|
|
self.diarizer = DiartWorker(config.diarization, config.sample_rate)
|
|
self.samples = np.empty(0, dtype=np.float32)
|
|
self.lock = threading.Lock()
|
|
self.emitted_until = 0.0
|
|
self.last_process_samples = 0
|
|
self.closed = False
|
|
|
|
def push_pcm(self, raw: bytes) -> None:
|
|
pcm = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
|
|
if pcm.size == 0:
|
|
return
|
|
self.diarizer.push(pcm)
|
|
with self.lock:
|
|
self.samples = np.concatenate([self.samples, pcm])
|
|
# Keep at most 20 minutes in memory. Long meetings are still recorded by the desktop app.
|
|
max_samples = self.config.sample_rate * 60 * 20
|
|
if self.samples.size > max_samples:
|
|
dropped = self.samples.size - max_samples
|
|
self.samples = self.samples[dropped:]
|
|
self.emitted_until = max(0.0, self.emitted_until - dropped / self.config.sample_rate)
|
|
|
|
def total_seconds(self) -> float:
|
|
with self.lock:
|
|
return self.samples.size / self.config.sample_rate
|
|
|
|
def should_process(self) -> bool:
|
|
with self.lock:
|
|
new_samples = self.samples.size - self.last_process_samples
|
|
return new_samples >= int(self.config.sample_rate * 1.4)
|
|
|
|
def transcribe_once(self) -> tuple[list[dict], Optional[str]]:
|
|
with self.lock:
|
|
total = self.samples.size
|
|
if total < self.config.sample_rate:
|
|
return [], None
|
|
window_size = min(total, self.config.sample_rate * 10)
|
|
audio = self.samples[-window_size:].copy()
|
|
self.last_process_samples = total
|
|
window_start = max(0.0, total / self.config.sample_rate - len(audio) / self.config.sample_rate)
|
|
|
|
segments_iter, _ = self.model.transcribe(
|
|
audio,
|
|
language=None if self.config.language in ("", "auto") else self.config.language,
|
|
beam_size=1,
|
|
best_of=1,
|
|
vad_filter=True,
|
|
word_timestamps=True,
|
|
condition_on_previous_text=False,
|
|
temperature=0.0,
|
|
)
|
|
segments = list(segments_iter)
|
|
results: list[dict] = []
|
|
partial: Optional[str] = None
|
|
safe_cutoff = total / self.config.sample_rate - 0.65
|
|
|
|
for seg in segments:
|
|
text = seg.text.strip()
|
|
if not text:
|
|
continue
|
|
abs_start = window_start + float(seg.start)
|
|
abs_end = window_start + float(seg.end)
|
|
if abs_end <= self.emitted_until + 0.08:
|
|
continue
|
|
if abs_end > safe_cutoff:
|
|
partial = text
|
|
continue
|
|
speaker_id = self.diarizer.resolve(abs_start, abs_end)
|
|
results.append({
|
|
"id": str(uuid.uuid4()),
|
|
"speakerId": speaker_id,
|
|
"start": abs_start,
|
|
"end": abs_end,
|
|
"text": text,
|
|
"final": True,
|
|
})
|
|
self.emitted_until = max(self.emitted_until, abs_end)
|
|
|
|
return results, partial
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
self.diarizer.close()
|
|
|
|
|
|
class MockSession:
|
|
def __init__(self, config: EngineConfig):
|
|
self.config = config
|
|
self.sample_count = 0
|
|
self.index = 0
|
|
self.lines = [
|
|
("speaker_1", "We should finalize the upload flow before Friday."),
|
|
("speaker_2", "I can take the local recording and video playback work."),
|
|
("speaker_1", "Great. Keep the backend in processing state for this demo."),
|
|
("speaker_3", "The overlay should keep updating while another application is open."),
|
|
]
|
|
|
|
def push_pcm(self, raw: bytes) -> Optional[dict]:
|
|
self.sample_count += len(raw) // 2
|
|
threshold = self.config.sample_rate * 2.4
|
|
if self.sample_count < threshold:
|
|
return None
|
|
self.sample_count = 0
|
|
speaker, text = self.lines[self.index % len(self.lines)]
|
|
self.index += 1
|
|
end = self.index * 2.4
|
|
return {
|
|
"id": str(uuid.uuid4()), "speakerId": speaker,
|
|
"start": end - 2.1, "end": end, "text": text, "final": True,
|
|
}
|
|
|
|
def close(self) -> None:
|
|
pass
|
|
|
|
|
|
app = FastAPI(title="MeetVault Local Speech Engine")
|
|
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
|
|
MOCK_ONLY = False
|
|
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return {"ok": True, "mode": "mock" if MOCK_ONLY else "real-capable", "sample_rate": SAMPLE_RATE}
|
|
|
|
|
|
@app.websocket("/ws/live")
|
|
async def live(websocket: WebSocket):
|
|
await websocket.accept()
|
|
session = None
|
|
process_task = None
|
|
try:
|
|
first = await websocket.receive_text()
|
|
payload = json.loads(first)
|
|
config = EngineConfig(
|
|
model=str(payload.get("model", "small")),
|
|
compute=str(payload.get("compute", "auto")),
|
|
language=str(payload.get("language", "auto")),
|
|
diarization=bool(payload.get("diarization", True)),
|
|
sample_rate=int(payload.get("sample_rate", SAMPLE_RATE)),
|
|
)
|
|
if config.sample_rate != SAMPLE_RATE:
|
|
await websocket.send_json({"type": "error", "message": "MeetVault speech engine currently requires 16 kHz mono PCM."})
|
|
return
|
|
|
|
if MOCK_ONLY:
|
|
session = MockSession(config)
|
|
await websocket.send_json({"type": "ready", "engine": "mock-whisper", "diarization": "mock-diart"})
|
|
else:
|
|
try:
|
|
session = await asyncio.to_thread(LiveSpeechSession, config)
|
|
except Exception as exc:
|
|
await websocket.send_json({
|
|
"type": "error",
|
|
"message": (
|
|
f"Could not initialize Whisper/diart: {exc}. "
|
|
"Use --mock for UI testing. Real diarization may require Hugging Face authentication and accepted pyannote model terms."
|
|
),
|
|
})
|
|
return
|
|
diarization_state = "diart ready" if session.diarizer.enabled else f"diart unavailable: {session.diarizer.error}"
|
|
await websocket.send_json({"type": "ready", "engine": f"Whisper {config.model}", "diarization": diarization_state})
|
|
|
|
async def processor():
|
|
while not session.closed:
|
|
await asyncio.sleep(0.35)
|
|
if not session.should_process():
|
|
continue
|
|
try:
|
|
finals, partial = await asyncio.to_thread(session.transcribe_once)
|
|
for segment in finals:
|
|
await websocket.send_json({"type": "final", "segment": segment})
|
|
if partial:
|
|
await websocket.send_json({"type": "partial", "text": partial})
|
|
except Exception as exc:
|
|
await websocket.send_json({"type": "error", "message": f"Live transcription error: {exc}"})
|
|
await asyncio.sleep(1)
|
|
|
|
process_task = asyncio.create_task(processor())
|
|
|
|
while True:
|
|
message = await websocket.receive()
|
|
if message.get("bytes") is not None:
|
|
raw = message["bytes"]
|
|
if isinstance(session, MockSession):
|
|
segment = session.push_pcm(raw)
|
|
if segment:
|
|
await websocket.send_json({"type": "final", "segment": segment})
|
|
else:
|
|
session.push_pcm(raw)
|
|
elif message.get("text"):
|
|
command = json.loads(message["text"])
|
|
if command.get("type") == "stop":
|
|
break
|
|
except WebSocketDisconnect:
|
|
pass
|
|
except Exception as exc:
|
|
try:
|
|
await websocket.send_json({"type": "error", "message": str(exc)})
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
if session:
|
|
session.close()
|
|
if process_task:
|
|
process_task.cancel()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--host", default="127.0.0.1")
|
|
parser.add_argument("--port", type=int, default=8765)
|
|
parser.add_argument("--mock", action="store_true", help="Run without Whisper/diart models")
|
|
args = parser.parse_args()
|
|
global MOCK_ONLY
|
|
MOCK_ONLY = args.mock
|
|
uvicorn.run(app, host=args.host, port=args.port, log_level="info")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|