feat: add WhisperX diarization and speaker transcript UI

- add WhisperX diarization support to the Whisper VM server
- normalize speaker timestamp segments from Whisper responses
- document Hugging Face/pyannote VM setup and health checks
- show diarized speaker transcript blocks in record and transcript views
- group consecutive segments from the same speaker
- remove duplicate paragraph transcript display when diarized segments exist
- let diarized transcript content expand without an inner scrollbar

scripts/orphion-whisper.service

@@ -8,9 +8,15 @@ Type=simple
 User=cezen
 Group=cezen
 WorkingDirectory=/home/cezen/whisper
+EnvironmentFile=-/home/cezen/whisper/.env
 Environment=WHISPER_MODEL=large-v3
 Environment=WHISPER_MODEL_DIR=/home/cezen/whisper/models
 Environment=WHISPER_FILE_FIELD=file
+Environment=WHISPERX_DIARIZATION=true
+Environment=WHISPERX_DEVICE=cuda
+Environment=WHISPERX_COMPUTE_TYPE=float16
+Environment=WHISPERX_BATCH_SIZE=8
+Environment=WHISPERX_DIARIZATION_MODEL=pyannote/speaker-diarization-community-1
 ExecStart=/home/cezen/whisper/venv/bin/python /home/cezen/whisper/server.py --host 0.0.0.0 --port 8000
 Restart=always
 RestartSec=5

scripts/whisper_http_server.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Tiny Faster-Whisper HTTP API for Orphion.
+"""Tiny Faster-Whisper/WhisperX HTTP API for Orphion.
 
 Endpoints:
   GET /health
@@ -10,6 +10,8 @@ from __future__ import annotations
 
 import argparse
 import cgi
+import gc
+import inspect
 import json
 import os
 import tempfile
@@ -21,6 +23,21 @@ from pathlib import Path
 
 from faster_whisper import WhisperModel
 
+try:
+    import torch
+except Exception:  # pragma: no cover - torch is optional for CPU fallback
+    torch = None
+
+try:
+    import whisperx
+except Exception:  # pragma: no cover - keeps existing Faster-Whisper path available
+    whisperx = None
+
+try:
+    import whisperx.diarize as whisperx_diarize
+except Exception:  # pragma: no cover - older/newer WhisperX layouts vary
+    whisperx_diarize = None
+
 try:
     import ctranslate2
 except Exception:  # pragma: no cover - only used for runtime device selection
@@ -29,6 +46,11 @@ except Exception:  # pragma: no cover - only used for runtime device selection
 
 MODEL_LOCK = threading.Lock()
 MODEL = None
+WHISPERX_LOCK = threading.Lock()
+INFERENCE_LOCK = threading.Lock()
+WHISPERX_MODEL = None
+ALIGN_MODELS = {}
+DIARIZATION_PIPELINE = None
 
 
 def choose_device() -> str:
@@ -50,7 +72,35 @@ def default_compute_type(device: str) -> str:
     configured = os.getenv("WHISPER_COMPUTE_TYPE", "").strip()
     if configured:
         return configured
-    return "float16" if device == "cuda" else "int8"
+    return "float16" if device.startswith("cuda") else "int8"
+
+
+def diarization_enabled() -> bool:
+    return os.getenv("WHISPERX_DIARIZATION", "true").lower() in {"1", "true", "yes", "on"}
+
+
+def whisperx_device() -> str:
+    configured = os.getenv("WHISPERX_DEVICE", "").strip()
+    if configured:
+        return configured
+    selected = choose_device()
+    return "cuda" if selected == "cuda" else selected
+
+
+def whisperx_compute_type(device: str) -> str:
+    configured = os.getenv("WHISPERX_COMPUTE_TYPE", "").strip()
+    if configured:
+        return configured
+    return "float16" if device.startswith("cuda") else default_compute_type(device)
+
+
+def huggingface_token() -> str | None:
+    return (
+        os.getenv("HUGGINGFACE_TOKEN")
+        or os.getenv("HF_TOKEN")
+        or os.getenv("PYANNOTE_AUTH_TOKEN")
+        or None
+    )
 
 
 def get_model() -> WhisperModel:
@@ -71,6 +121,196 @@ def get_model() -> WhisperModel:
     return MODEL
 
 
+def get_whisperx_model():
+    global WHISPERX_MODEL
+    if WHISPERX_MODEL is not None:
+        return WHISPERX_MODEL
+
+    if whisperx is None:
+        raise RuntimeError("whisperx is not installed")
+
+    with WHISPERX_LOCK:
+        if WHISPERX_MODEL is None:
+            device = whisperx_device()
+            model_name = os.getenv("WHISPERX_MODEL", os.getenv("WHISPER_MODEL", "large-v3"))
+            kwargs = {
+                "compute_type": whisperx_compute_type(device),
+                "download_root": os.getenv("WHISPER_MODEL_DIR") or None,
+            }
+            try:
+                WHISPERX_MODEL = whisperx.load_model(model_name, device, **kwargs)
+            except TypeError:
+                kwargs.pop("download_root", None)
+                WHISPERX_MODEL = whisperx.load_model(model_name, device, **kwargs)
+    return WHISPERX_MODEL
+
+
+def get_align_model(language_code: str | None, device: str):
+    if whisperx is None:
+        raise RuntimeError("whisperx is not installed")
+
+    language = language_code or os.getenv("WHISPERX_ALIGN_LANGUAGE", "en")
+    key = (language, device)
+    if key not in ALIGN_MODELS:
+        ALIGN_MODELS[key] = whisperx.load_align_model(language_code=language, device=device)
+    return ALIGN_MODELS[key]
+
+
+def get_diarization_pipeline(device: str):
+    global DIARIZATION_PIPELINE
+    if whisperx is None:
+        raise RuntimeError("whisperx is not installed")
+    if DIARIZATION_PIPELINE is not None:
+        return DIARIZATION_PIPELINE
+
+    token = huggingface_token()
+    if not token:
+        raise RuntimeError("HuggingFace token is required for speaker diarization")
+
+    with WHISPERX_LOCK:
+        if DIARIZATION_PIPELINE is None:
+            pipeline_factory = getattr(whisperx, "DiarizationPipeline", None)
+            if pipeline_factory is None and whisperx_diarize is not None:
+                pipeline_factory = getattr(whisperx_diarize, "DiarizationPipeline", None)
+            if pipeline_factory is None:
+                raise RuntimeError("WhisperX diarization pipeline is not available")
+            parameters = inspect.signature(pipeline_factory).parameters
+            if "use_auth_token" in parameters:
+                DIARIZATION_PIPELINE = pipeline_factory(use_auth_token=token, device=device)
+            elif "auth_token" in parameters:
+                DIARIZATION_PIPELINE = pipeline_factory(auth_token=token, device=device)
+            elif "token" in parameters:
+                DIARIZATION_PIPELINE = pipeline_factory(
+                    model_name=os.getenv("WHISPERX_DIARIZATION_MODEL") or None,
+                    token=token,
+                    device=device,
+                )
+            else:
+                DIARIZATION_PIPELINE = pipeline_factory(device=device)
+    return DIARIZATION_PIPELINE
+
+
+def cleanup_gpu_memory() -> None:
+    gc.collect()
+    if torch is not None and getattr(torch, "cuda", None) is not None:
+        try:
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        except Exception:
+            pass
+
+
+def normalize_speaker(raw: str | None, mapping: dict[str, str]) -> str:
+    if not raw:
+        return "Speaker 1"
+    key = str(raw)
+    if key not in mapping:
+        mapping[key] = f"Speaker {len(mapping) + 1}"
+    return mapping[key]
+
+
+def segment_speaker(segment: dict) -> str | None:
+    if segment.get("speaker"):
+        return segment.get("speaker")
+    counts: dict[str, int] = {}
+    for word in segment.get("words") or []:
+        speaker = word.get("speaker")
+        if speaker:
+            counts[speaker] = counts.get(speaker, 0) + 1
+    if not counts:
+        return None
+    return max(counts.items(), key=lambda item: item[1])[0]
+
+
+def normalize_segments(segments: list[dict], include_speakers: bool = True) -> list[dict]:
+    speaker_map: dict[str, str] = {}
+    normalized = []
+    for segment in segments:
+        text = str(segment.get("text") or "").strip()
+        if not text:
+            continue
+        item = {
+            "start": round(float(segment.get("start") or 0), 3),
+            "end": round(float(segment.get("end") or 0), 3),
+            "text": text,
+        }
+        if include_speakers:
+            item["speaker"] = normalize_speaker(segment_speaker(segment), speaker_map)
+        normalized.append(item)
+    return normalized
+
+
+def transcribe_with_whisperx(audio_path: str) -> dict:
+    device = whisperx_device()
+    model = get_whisperx_model()
+    audio = whisperx.load_audio(audio_path)
+    batch_size = int(os.getenv("WHISPERX_BATCH_SIZE", os.getenv("WHISPER_BATCH_SIZE", "8")))
+    result = model.transcribe(audio, batch_size=batch_size)
+
+    language = result.get("language")
+    try:
+        align_model, metadata = get_align_model(language, device)
+        result = whisperx.align(
+            result.get("segments", []),
+            align_model,
+            metadata,
+            audio,
+            device,
+            return_char_alignments=False,
+        )
+    except Exception:
+        traceback.print_exc()
+
+    diarization_failed = False
+    try:
+        diarize_model = get_diarization_pipeline(device)
+        min_speakers = os.getenv("WHISPERX_MIN_SPEAKERS")
+        max_speakers = os.getenv("WHISPERX_MAX_SPEAKERS")
+        kwargs = {}
+        if min_speakers:
+            kwargs["min_speakers"] = int(min_speakers)
+        if max_speakers:
+            kwargs["max_speakers"] = int(max_speakers)
+        diarize_segments = diarize_model(audio, **kwargs)
+        result = whisperx.assign_word_speakers(diarize_segments, result)
+    except Exception:
+        diarization_failed = True
+        traceback.print_exc()
+
+    timestamps = normalize_segments(result.get("segments", []), include_speakers=not diarization_failed)
+    transcript_text = " ".join(segment["text"] for segment in timestamps if segment.get("text"))
+    return {
+        "transcript_text": transcript_text,
+        "language": language,
+        "duration": float(audio.shape[0]) / 16000 if hasattr(audio, "shape") else None,
+        "timestamps": timestamps,
+        "diarization": "fallback" if diarization_failed else "completed",
+    }
+
+
+def transcribe_with_faster_whisper(audio_path: str) -> dict:
+    model = get_model()
+    segments, info = model.transcribe(
+        audio_path,
+        beam_size=int(os.getenv("WHISPER_BEAM_SIZE", "5")),
+        vad_filter=os.getenv("WHISPER_VAD_FILTER", "true").lower() in {"1", "true", "yes", "on"},
+    )
+    timestamps = []
+    transcript_parts = []
+    for segment in segments:
+        text = segment.text.strip()
+        transcript_parts.append(text)
+        timestamps.append({"start": segment.start, "end": segment.end, "text": text})
+
+    return {
+        "transcript_text": " ".join(part for part in transcript_parts if part),
+        "language": getattr(info, "language", None),
+        "duration": getattr(info, "duration", None),
+        "timestamps": timestamps,
+        "diarization": "disabled",
+    }
+
+
 class WhisperHandler(BaseHTTPRequestHandler):
     server_version = "OrphionWhisper/1.0"
 
@@ -84,6 +324,11 @@ class WhisperHandler(BaseHTTPRequestHandler):
                 "status": "ok",
                 "model": os.getenv("WHISPER_MODEL", "large-v3"),
                 "device": choose_device(),
+                "whisperx": whisperx is not None,
+                "diarization": diarization_enabled(),
+                "diarization_ready": whisperx is not None
+                and diarization_enabled()
+                and bool(huggingface_token()),
             },
         )
 
@@ -129,28 +374,20 @@ class WhisperHandler(BaseHTTPRequestHandler):
                 temp_file.write(chunk)
 
         try:
-            model = get_model()
-            segments, info = model.transcribe(
-                temp_path,
-                beam_size=int(os.getenv("WHISPER_BEAM_SIZE", "5")),
-                vad_filter=os.getenv("WHISPER_VAD_FILTER", "true").lower() in {"1", "true", "yes", "on"},
-            )
-            timestamps = []
-            transcript_parts = []
-            for segment in segments:
-                text = segment.text.strip()
-                transcript_parts.append(text)
-                timestamps.append({"start": segment.start, "end": segment.end, "text": text})
+            with INFERENCE_LOCK:
+                try:
+                    if diarization_enabled():
+                        payload = transcribe_with_whisperx(temp_path)
+                    else:
+                        payload = transcribe_with_faster_whisper(temp_path)
+                except Exception:
+                    traceback.print_exc()
+                    payload = transcribe_with_faster_whisper(temp_path)
+                    payload["diarization"] = "fallback"
 
-            self.send_json(
-                {
-                    "transcript_text": " ".join(part for part in transcript_parts if part),
-                    "language": getattr(info, "language", None),
-                    "duration": getattr(info, "duration", None),
-                    "timestamps": timestamps,
-                },
-            )
+            self.send_json(payload)
         finally:
+            cleanup_gpu_memory()
             try:
                 os.unlink(temp_path)
             except OSError: