Update app.py

app.py

@@ -6,6 +6,7 @@ from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.middleware.cors import CORSMiddleware
 from faster_whisper import WhisperModel
 from df.enhance import enhance, init_df, load_audio, save_audio
+from pyannote.audio import Pipeline as DiarizationPipeline
 from pipeline.postprocess import process
 
 app = FastAPI()
@@ -34,6 +35,13 @@ align_labels = align_bundle.get_labels()
 align_dictionary = {c: i for i, c in enumerate(align_labels)}
 print("Word-alignment model ready!")
 
+print("Loading speaker-diarization model...", flush=True)
+HF_TOKEN = os.environ.get("HF_TOKEN")
+diarize_pipeline = DiarizationPipeline.from_pretrained(
+    "pyannote/speaker-diarization-3.1", use_auth_token=HF_TOKEN
+)
+print("Speaker-diarization model ready!")
+
 
 def realign_word_timestamps(waveform_16k, words, whisper_start, whisper_end):
     """
@@ -72,7 +80,6 @@ def realign_word_timestamps(waveform_16k, words, whisper_start, whisper_end):
         spans = torchaudio.functional.merge_tokens(aligned_tokens[0], scores[0])
 
         ratio = clip.shape[1] / emission.shape[1] / 16000
-        char_idx = 0
         new_words = []
         span_idx = 0
         for w, tok in clean_words:
@@ -92,6 +99,17 @@ def realign_word_timestamps(waveform_16k, words, whisper_start, whisper_end):
         return words
 
 
+def get_speaker_at(diarization, time_sec):
+    """
+    برای یه لحظه‌ی مشخص از صدا، می‌گه اون لحظه مال کدوم گوینده‌ست (SPEAKER_00, SPEAKER_01, ...).
+    اگه هیچ گوینده‌ای پیدا نشد، None برمی‌گردونه.
+    """
+    for turn, _, speaker in diarization.itertracks(yield_label=True):
+        if turn.start <= time_sec <= turn.end:
+            return speaker
+    return None
+
+
 @app.get("/health")
 async def health():
     return {"status": "ok", "message": "SubFlow backend is running"}
@@ -155,6 +173,9 @@ async def transcribe(websocket: WebSocket):
         if align_waveform.shape[0] > 1:
             align_waveform = align_waveform.mean(dim=0, keepdim=True)
 
+        print("Step 5.6: running speaker diarization")
+        diarization = diarize_pipeline(audio_to_transcribe)
+
         segments = []
         for seg in segments_raw:
             words = list(seg.words)
@@ -166,16 +187,20 @@ async def transcribe(websocket: WebSocket):
                 align_waveform, word_dicts, word_dicts[0]["start"], word_dicts[-1]["end"]
             )
 
+            mid_time = (realigned[0]["start"] + realigned[-1]["end"]) / 2
+            speaker = get_speaker_at(diarization, mid_time)
+
             segments.append({
                 "text": seg.text.strip(),
                 "start": realigned[0]["start"],
                 "end": realigned[-1]["end"],
                 "words": realigned,
+                "speaker": speaker,
             })
 
         print(f"Step 6: got {len(segments)} segments")
         for i, seg in enumerate(segments[:5]):
-            print(f"  seg[{i}]: '{seg['text']}' | words={len(seg['words'])} | dur={seg['end']-seg['start']:.2f}s")
+            print(f"  seg[{i}]: '{seg['text']}' | speaker={seg['speaker']} | dur={seg['end']-seg['start']:.2f}s")
 
         await websocket.send_json({"status": "processing"})
 
@@ -189,8 +214,6 @@ async def transcribe(websocket: WebSocket):
             "duration": info.duration,
         })
         print("Step 8: done sent!")
-
-        # کمی صبر کن تا client پیام رو دریافت کنه
         import asyncio
         await asyncio.sleep(1)