diarization1Mæló

app.py

@@ -1,82 +1,121 @@
 import os
 import gradio as gr
 import spaces
-import tempfile
+import webrtcvad
+import numpy as np
+from pydub import AudioSegment
+from sklearn.cluster import AgglomerativeClustering
+from transformers import pipeline, Wav2Vec2Processor, Wav2Vec2Model
 import torch
+import tempfile
 
-from transformers import pipeline
-from pyannote.audio import Pipeline
+ASR_MODEL = "palli23/whisper-small-sam_spjall"
 
-# ==========================================================
-# ZeroGPU SAFE GLOBALS FIX — PYANNOTE 3.1 CHECKPOINT COMPAT
-# ==========================================================
-from torch.serialization import add_safe_globals
-from pyannote.audio.core.task import Specifications
-from pyannote.audio.core.model import Model
+# Load speech embedding model (ECAPA)
+EMB_MODEL = "speechbrain/spkrec-ecapa-voxceleb"
+processor = Wav2Vec2Processor.from_pretrained(EMB_MODEL)
+embedder = Wav2Vec2Model.from_pretrained(EMB_MODEL)
 
-add_safe_globals({
-    "Specifications": Specifications,
-    "pyannote.audio.core.task.Specifications": Specifications,
-    "Model": Model,
-    "pyannote.audio.core.model.Model": Model,
-})
 
+def audio_to_frames(path, frame_ms=30):
+    audio = AudioSegment.from_file(path).set_channels(1).set_frame_rate(16000)
+    samples = np.array(audio.get_array_of_samples()).astype(np.int16)
+    frame_len = int(16000 * frame_ms / 1000)
+    for i in range(0, len(samples), frame_len):
+        yield samples[i:i + frame_len]
 
-ASR_MODEL = "palli23/whisper-small-sam_spjall"
-DIAR_MODEL = "pyannote/speaker-diarization-3.1"
 
+def extract_segments(path):
+    vad = webrtcvad.Vad(2)
+    frames = list(audio_to_frames(path))
 
-@spaces.GPU(duration=120)
-def transcribe_with_diarization(audio_path):
+    segments = []
+    current = []
 
-    if not audio_path:
-        return "Hladdu upp hljóðskrá."
+    for frame in frames:
+        if len(frame) < 480:
+            continue
+
+        is_speech = vad.is_speech(frame.tobytes(), 16000)
+        if is_speech:
+            current.append(frame)
+        else:
+            if current:
+                segments.append(np.concatenate(current))
+                current = []
+
+    if current:
+        segments.append(np.concatenate(current))
+
+    return segments
 
-    # ----------------------------
-    # Load diarization pipeline
-    # (NO token argument!)
-    # ----------------------------
-    diarization = Pipeline.from_pretrained(
-        DIAR_MODEL,
-        cache_dir="/home/user/.cache"
-    ).to("cuda")
 
-    diar = diarization(audio_path)
+def embed_audio(segment):
+    with torch.no_grad():
+        inputs = processor(segment, sampling_rate=16000, return_tensors="pt")
+        emb = embedder(**inputs).last_hidden_state.mean(dim=1)
+    return emb[0].numpy()
 
-    # ----------------------------
-    # Whisper ASR
-    # ----------------------------
-    asr = pipeline(
-        task="automatic-speech-recognition",
-        model=ASR_MODEL,
-        device=0
-    )
 
-    output_lines = []
+def cluster_speakers(embeddings, max_speakers=5):
+    X = np.stack(embeddings)
+    clustering = AgglomerativeClustering(
+        n_clusters=None,
+        distance_threshold=1.0
+    ).fit(X)
 
-    for turn, _, speaker in diar.itertracks(yield_label=True):
+    return clustering.labels_
 
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
-            diar.crop(audio_path, turn).export(tmp.name, format="wav")
-            seg_file = tmp.name
 
-        text = asr(seg_file)["text"].strip()
-        output_lines.append(f"[MÆLENDI {speaker}] {text}")
+asr = pipeline("automatic-speech-recognition",
+               model=ASR_MODEL, device=0)
 
-        os.unlink(seg_file)
 
-    return "\n".join(output_lines) or "Enginn texti fannst."
+@spaces.GPU(duration=120)
+def diarize_and_transcribe(audio_path):
+    if not audio_path:
+        return "Hladdu upp hljóðskrá"
+
+    # --- STEP 1: VAD speech detection ---
+    segments = extract_segments(audio_path)
+    if not segments:
+        return "Engin tala heyrðist í skránni."
+
+    embeddings = [embed_audio(seg) for seg in segments]
+
+    # --- STEP 2: Speaker clustering ---
+    labels = cluster_speakers(embeddings)
 
+    # --- STEP 3: ASR á hverju segmenti ---
+    out = []
+    for seg, spk in zip(segments, labels):
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+            audio = (seg.astype(np.int16)).tobytes()
+            temp_audio = AudioSegment(
+                data=audio,
+                sample_width=2,
+                frame_rate=16000,
+                channels=1
+            )
+            temp_audio.export(f.name, format="wav")
+            seg_path = f.name
 
-# ==========================================================
-# UI
-# ==========================================================
+        txt = asr(seg_path)["text"].strip()
+        out.append(f"[MÆLENDI {spk}] {txt}")
+        os.unlink(seg_path)
+
+    return "\n".join(out)
+
+
+# --- Gradio UI ---
 with gr.Blocks() as demo:
-    gr.Markdown("# ���️ Íslenskt ASR + mælendagreining (ZeroGPU)")
-    audio = gr.Audio(type="filepath", label="Hlaða inn hljóði (.wav / .mp3)")
-    out = gr.Textbox(lines=25, label="Útskrift")
+    gr.Markdown("# Íslenskt ASR + VAD mælendagreining (WebRTC)")
+    gr.Markdown("Virkar á ZeroGPU\nHladdu upp .mp3 / .wav (allt að 5 mín)")
+
+    audio = gr.Audio(type="filepath")
+    btn = gr.Button("Transcribe með mælendum")
+    out = gr.Textbox(lines=35)
 
-    btn = gr.Button("Transcribe")
-    btn.click(transcribe_with_diarization, inputs=audio, outputs=out)
+    btn.click(diarize_and_transcribe, inputs=audio, outputs=out)
 
 demo.launch(auth=("beta", "beta2025"))

requirements.txt

@@ -1,7 +1,7 @@
-gradio
-transformers
-torch
-spaces
-pyannote.audio
-librosa
-soundfile
+torch==2.0.1
+transformers==4.40.2
+webrtcvad
+pydub
+numpy
+scikit-learn
+sentencepiece