Update app.py

app.py

@@ -1,117 +1,54 @@
 import gradio as gr
-import torch
-import numpy as np
-import librosa
+import os
+from pyannote.audio import Pipeline
 
-from speechbrain.pretrained import EncoderClassifier
-from sklearn.cluster import AgglomerativeClustering
-from sklearn.metrics.pairwise import cosine_similarity
-
-# 🔥 Strong embedding model (better than resemblyzer)
-classifier = EncoderClassifier.from_hparams(
-    source="speechbrain/spkrec-ecapa-voxceleb",
-    run_opts={"device": "cpu"}
+# Load pipeline once (important for speed)
+pipeline = Pipeline.from_pretrained(
+    "pyannote/speaker-diarization",
+    use_auth_token=os.getenv("HF_TOKEN")
 )
 
-def diarize(audio_path):
-    wav, sr = librosa.load(audio_path, sr=16000)
-
-    # 🔥 VAD (simple energy-based)
-    frame_length = int(0.025 * sr)
-    hop_length = int(0.01 * sr)
+# Optional: use GPU if available
+try:
+    import torch
+    if torch.cuda.is_available():
+        pipeline.to(torch.device("cuda"))
+except:
+    pass
 
-    energy = np.array([
-        np.sum(np.abs(wav[i:i+frame_length]**2))
-        for i in range(0, len(wav), hop_length)
-    ])
 
-    threshold = np.mean(energy) * 0.5
-    speech_frames = energy > threshold
+def diarize(audio_path):
+    diarization = pipeline(audio_path)
 
-    # Convert frames to segments
     segments = []
-    start = None
-
-    for i, val in enumerate(speech_frames):
-        if val and start is None:
-            start = i * hop_length / sr
-        elif not val and start is not None:
-            end = i * hop_length / sr
-            segments.append((start, end))
-            start = None
-
-    if start is not None:
-        segments.append((start, len(wav)/sr))
-
-    # 🔥 Extract embeddings
-    embeddings = []
-    valid_segments = []
-
-    for (s, e) in segments:
-        if e - s < 0.5:
-            continue
-
-        audio_chunk = wav[int(s*sr):int(e*sr)]
-        emb = classifier.encode_batch(torch.tensor(audio_chunk).unsqueeze(0))
-        emb = emb.squeeze().detach().numpy()
-
-        embeddings.append(emb)
-        valid_segments.append((s, e))
-
-    if len(embeddings) == 0:
-        return {"segments": []}
-
-    embeddings = np.array(embeddings)
-
-    # 🔥 Clustering
-    sim_matrix = cosine_similarity(embeddings)
-
-    clustering = AgglomerativeClustering(
-        n_clusters=None,
-        distance_threshold=0.7
-    )
-
-    labels = clustering.fit_predict(embeddings)
-
-    # 🔥 Merge same speaker segments
-    merged = []
-    prev_label = labels[0]
-    start = valid_segments[0][0]
-    end = valid_segments[0][1]
+    speaker_map = {}
 
-    for i in range(1, len(labels)):
-        if labels[i] == prev_label:
-            end = valid_segments[i][1]
-        else:
-            merged.append({
-                "speaker": f"Speaker {prev_label+1}",
-                "start": round(start, 2),
-                "end": round(end, 2)
-            })
-            prev_label = labels[i]
-            start = valid_segments[i][0]
-            end = valid_segments[i][1]
+    for turn, _, speaker in diarization.itertracks(yield_label=True):
+        if speaker not in speaker_map:
+            speaker_map[speaker] = f"Speaker {len(speaker_map) + 1}"
 
-    merged.append({
-        "speaker": f"Speaker {prev_label+1}",
-        "start": round(start, 2),
-        "end": round(end, 2)
-    })
+        segments.append({
+            "speaker": speaker_map[speaker],
+            "start": round(turn.start, 2),
+            "end": round(turn.end, 2)
+        })
 
-    return {"segments": merged}
+    return {"segments": segments}
 
 
 def process(audio):
     if audio is None:
-        return {"error": "No audio"}
+        return {"error": "No audio provided"}
     return diarize(audio)
 
 
+# Gradio Interface
 demo = gr.Interface(
     fn=process,
     inputs=gr.Audio(sources=["upload"], type="filepath"),
     outputs=gr.JSON(),
-    title="High Accuracy Speaker Diarization (No Token)"
+    title="Speaker Diarization API (pyannote)",
+    description="High accuracy speaker labeling with timestamps"
 )
 
 if __name__ == "__main__":