Add speaker diarization

app.py

@@ -1,9 +1,12 @@
 import os
 import re
+import tempfile
 
 import torch
 import gradio as gr
 from transformers import pipeline
+from pydub import AudioSegment
+from pyannote.audio import Pipeline as DiarizationPipeline
 
 import spaces  # zeroGPU support
 from funasr import AutoModel
@@ -67,10 +70,10 @@ SENSEVOICE_LANGUAGES = ["auto", "zh", "yue", "en", "ja", "ko", "nospeech"]
 # —————— Caches ——————
 whisper_pipes = {}
 sense_models = {}
+dar_pipe = None
 
 # —————— Helpers ——————
 def get_whisper_pipe(model_id: str, device: int):
-    # HuggingFace pipeline caching by model and device (-1=cpu, 0=gpu)
     key = (model_id, device)
     if key not in whisper_pipes:
         whisper_pipes[key] = pipeline(
@@ -96,32 +99,50 @@ def get_sense_model(model_id: str):
         )
     return sense_models[model_id]
 
-# —————— Whisper Transcribers ——————
-@spaces.GPU
-def _transcribe_whisper_gpu(model_id: str, language: str, audio_path: str):
-    pipe = get_whisper_pipe(model_id, device=0)
-    if language == "auto":
-        result = pipe(audio_path)
-    else:
-        result = pipe(audio_path, generate_kwargs={"language": language})
-    return result.get("text", "").strip()
 
-def _transcribe_whisper_cpu(model_id: str, language: str, audio_path: str):
-    pipe = get_whisper_pipe(model_id, device=-1)
+def get_diarization_pipe():
+    global dar_pipe
+    if dar_pipe is None:
+        dar_pipe = DiarizationPipeline.from_pretrained(
+            "pyannote/speaker-diarization@2.1",
+            use_auth_token=True
+        )
+    return dar_pipe
+
+# —————— Transcription Functions ——————
+def transcribe_whisper(model_id: str, language: str, audio_path: str, device_sel: str, enable_diar: bool):
+    # select device
+    use_gpu = (device_sel == "GPU" and torch.cuda.is_available())
+    device = 0 if use_gpu else -1
+    pipe = get_whisper_pipe(model_id, device)
+    # full transcription
     if language == "auto":
         result = pipe(audio_path)
     else:
         result = pipe(audio_path, generate_kwargs={"language": language})
-    return result.get("text", "").strip()
-
-# Wrapper selects GPU or CPU path
-def transcribe_whisper(model_id: str, language: str, audio_path: str, device_sel: str):
-    if device_sel == "GPU" and torch.cuda.is_available():
-        return _transcribe_whisper_gpu(model_id, language, audio_path)
-    else:
-        return _transcribe_whisper_cpu(model_id, language, audio_path)
+    transcript = result.get("text", "").strip()
+    diar_text = ""
+    # optional diarization
+    if enable_diar:
+        diarizer = get_diarization_pipe()
+        diarization = diarizer(audio_path)
+        snippets = []
+        for turn, _, speaker in diarization.itertracks(yield_label=True):
+            start_ms = int(turn.start * 1000)
+            end_ms = int(turn.end * 1000)
+            segment = AudioSegment.from_file(audio_path)[start_ms:end_ms]
+            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+                segment.export(tmp.name, format="wav")
+                if language == "auto":
+                    seg_out = pipe(tmp.name)
+                else:
+                    seg_out = pipe(tmp.name, generate_kwargs={"language": language})
+            os.unlink(tmp.name)
+            txt = seg_out.get("text", "").strip()
+            snippets.append(f"[{speaker}] {txt}")
+        diar_text = "\n".join(snippets)
+    return transcript, diar_text
 
-# —————— SenseVoice Transcriber ——————
 @spaces.GPU
 def transcribe_sense(model_id: str, language: str, audio_path: str, enable_punct: bool):
     model = get_sense_model(model_id)
@@ -142,7 +163,7 @@ def transcribe_sense(model_id: str, language: str, audio_path: str, enable_punct
 # —————— Gradio UI ——————
 demo = gr.Blocks()
 with demo:
-    gr.Markdown("## Whisper vs. SenseVoice Transcription (Language & Device)")
+    gr.Markdown("## Whisper vs. SenseVoice Transcription (with Language, Device & Diarization)")
 
     audio_input = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Audio Input")
 
@@ -153,12 +174,14 @@ with demo:
             whisper_dd = gr.Dropdown(choices=WHISPER_MODELS, value=WHISPER_MODELS[0], label="Whisper Model")
             whisper_lang = gr.Dropdown(choices=WHISPER_LANGUAGES, value="auto", label="Whisper Language")
             device_radio = gr.Radio(choices=["GPU", "CPU"], value="GPU", label="Device")
+            diar_check = gr.Checkbox(label="Enable Speaker Diarization", value=False)
             whisper_btn = gr.Button("Transcribe with Whisper")
             out_whisper = gr.Textbox(label="Whisper Transcript")
+            out_diar = gr.Textbox(label="Diarized Transcript (Whisper)")
             whisper_btn.click(
                 fn=transcribe_whisper,
-                inputs=[whisper_dd, whisper_lang, audio_input, device_radio],
-                outputs=[out_whisper]
+                inputs=[whisper_dd, whisper_lang, audio_input, device_radio, diar_check],
+                outputs=[out_whisper, out_diar]
             )
 
         # SenseVoice column

requirements.txt

@@ -1,12 +1,20 @@
-gradio>=5.0
-transformers>=4.30.0
-torch>=1.10
-torchaudio>=0.10
-accelerate>=0.20.0
-ffmpeg-python>=0.2.0
-
-# FunASR core & post-processing
-funasr>=0.1.1
-
-# HF utilities (already present on Spaces, but pin if you need a specific version)
-huggingface-hub>=0.14.1
+# Gradio UI
+gradio>=3.39.0
+
+# Core ASR
+torch>=2.0.0
+transformers>=4.35.0
+
+# FunASR SenseVoice
+funasr>=0.6.4
+
+# Audio handling
+pydub>=0.25.1
+ffmpeg-python>=0.2.0   # wrapper for ffmpeg; you’ll still need system ffmpeg installed
+
+# Speaker Diarization
+pyannote.audio>=2.1.1
+huggingface-hub>=0.18.0  # for pyannote model download/auth
+
+# (Optional) if you want GPU‐accelerated pipelines outside of HF Spaces
+# accelerate>=0.20.0