From 9a93cb6ffde0444ab48e11d43cde520282d682a8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?IT=E7=8B=97?= <[email protected]>
Date: Wed, 12 Aug 2026 01:34:45 +0800
Subject: [PATCH] transcribe_server: bypass whisperx DiarizationPipeline, use
 raw pyannote Pipeline
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Speaker labels were empty because whisperx 3.8.5's DiarizationPipeline wrapper
has a compatibility bug with pyannote.audio 4.0.4: the wrapper silently returns
empty segments (the 'fallback to plain text' message in 3.8.5 logs), so
segments.speaker becomes '' instead of 'SPEAKER_XX'.

Fix: load pyannote.audio.Pipeline directly and convert DiarizeOutput →
DataFrame the same way whisperx/diarize.py does internally. Verified with 30s
Cantonese clip from real recording (chunk_0.mp3):

  speaker: SPEAKER_00  (was '')
  text: '那就係F5嘅Low Balancer用喺你哋嘅...特別感謝Sherry...'
  took: 34s for 30s audio (1.13x realtime, CPU int8 + diarize=1)
  curl http://localhost:8765/transcribe -F file=@test.wav -F language=yue -F diarize=1
---
 transcribe_server.py | 40 ++++++++++++++++++++++++++++------------
 1 file changed, 28 insertions(+), 12 deletions(-)

diff --git a/transcribe_server.py b/transcribe_server.py
index 8635209..cece0ba 100644
--- a/transcribe_server.py
+++ b/transcribe_server.py
@@ -47,9 +47,18 @@ def get_model():
 def get_diarize():
     global _diarize
     if _diarize is None:
-        from whisperx.diarize import DiarizationPipeline
-        print("Loading diarization model...")
-        _diarize = DiarizationPipeline(token=os.environ["HF_TOKEN"], device="cpu", model_name="pyannote/speaker-diarization-3.1")
+        # BUGFIX (2026-08-12): whisperx 3.8.5's DiarizationPipeline wrapper has a bug
+        # with pyannote.audio 4.0.4 (returns empty segs, speaker=""). Use raw pyannote
+        # Pipeline directly — confirmed working: 60s audio → 1 seg with SPEAKER_00.
+        from pyannote.audio import Pipeline
+        import torch
+        print("Loading diarization model (raw pyannote Pipeline)...")
+        pipeline = Pipeline.from_pretrained(
+            "pyannote/speaker-diarization-3.1",
+            token=os.environ["HF_TOKEN"],
+        )
+        pipeline.to(torch.device("cpu"))
+        _diarize = pipeline
         print("Diarization loaded!")
     return _diarize
 
@@ -103,15 +112,22 @@ def _do_transcribe(audio_path: str, language: str, do_diarize: bool):
     if do_diarize:
         try:
             import torch
-            diarize_model = get_diarize()
-            # BUGFIX (2026-08-11): pyannote's torchcodec can't load (FFmpeg 8 only has
-            # libavutil.60, venv needs .59). Pre-load audio as tensor dict to bypass.
-            if isinstance(audio, np.ndarray):
-                audio_input = {"waveform": torch.from_numpy(audio).unsqueeze(0).float(),
-                               "sample_rate": 16000}
-                diar_segments = diarize_model(audio_input)
-            else:
-                diar_segments = diarize_model(audio)
+            import pandas as pd
+            diarize_pipeline = get_diarize()
+            # BUGFIX (2026-08-12): raw pyannote Pipeline; bypass whisperx wrapper.
+            # Pass dict {waveform, sample_rate} — same shape whisperx used internally.
+            audio_input = {"waveform": torch.from_numpy(audio).unsqueeze(0).float(),
+                           "sample_rate": 16000}
+            diarize_output = diarize_pipeline(audio_input, min_speakers=1, max_speakers=4)
+            # Raw pyannote returns DiarizeOutput; convert Annotation → DataFrame
+            # the same way whisperx/diarize.py does.
+            diarization = diarize_output.speaker_diarization
+            diar_segments = pd.DataFrame(
+                diarization.itertracks(yield_label=True),
+                columns=['segment', 'label', 'speaker'],
+            )
+            diar_segments['start'] = diar_segments['segment'].apply(lambda x: x.start)
+            diar_segments['end'] = diar_segments['segment'].apply(lambda x: x.end)
             result = whisperx.assign_word_speakers(diar_segments, result)
             for seg in result["segments"]:
                 segments_with_speakers.append({