Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import tempfile
-import numpy as np
 import torch
+import torchaudio
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.middleware.cors import CORSMiddleware
 from faster_whisper import WhisperModel
@@ -25,7 +25,72 @@ print("Loading DeepFilterNet model...", flush=True)
 df_model, df_state, _ = init_df()
 print("DeepFilterNet model ready!")
 
-USE_DEEPFILTERNET = False  # برای تست بدون DeepFilterNet، این رو موقتاً False کن
+USE_DEEPFILTERNET = False  # غیرفعال - شواهد نشون داد کمکی نمی‌کنه و بعضی صداها رو خراب می‌کنه
+
+print("Loading word-alignment model...", flush=True)
+align_bundle = torchaudio.pipelines.MMS_FA
+align_model = align_bundle.get_model().to("cpu")
+align_labels = align_bundle.get_labels()
+align_dictionary = {c: i for i, c in enumerate(align_labels)}
+print("Word-alignment model ready!")
+
+
+def realign_word_timestamps(waveform_16k, words, whisper_start, whisper_end):
+    """
+    برای یه جمله، به‌جای زمان‌بندی حدسی خودِ Whisper، با یه مدل تخصصی جدا
+    (wav2vec2 از طریق torchaudio) دقیقاً چک می‌کنه هر کلمه از کجای صدا گفته شده.
+    اگه به هر دلیلی هم‌ترازی جواب نداد، همون زمان‌بندی قبلی Whisper دست‌نخورده می‌مونه.
+    """
+    try:
+        margin = 0.5
+        start_sample = max(0, int((whisper_start - margin) * 16000))
+        end_sample = int((whisper_end + margin) * 16000)
+        clip = waveform_16k[:, start_sample:end_sample]
+
+        clean_words = []
+        for w in words:
+            token = "".join(ch for ch in w["word"].upper() if ch in align_dictionary)
+            if token:
+                clean_words.append((w, token))
+
+        if not clean_words:
+            return words
+
+        transcript = "".join(tok for _, tok in clean_words)
+        tokens = [align_dictionary[c] for c in transcript]
+
+        with torch.inference_mode():
+            emission, _ = align_model(clip)
+
+        targets = torch.tensor([tokens], dtype=torch.int32)
+        input_lengths = torch.tensor([emission.shape[1]])
+        target_lengths = torch.tensor([len(tokens)])
+
+        aligned_tokens, scores = torchaudio.functional.forced_align(
+            emission, targets, input_lengths, target_lengths, blank=0
+        )
+        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:
+            n_chars = len(tok)
+            matched = [s for s in spans[span_idx:] if s.token != 0][:n_chars]
+            if len(matched) < n_chars:
+                new_words.append(w)
+                continue
+            w_start = matched[0].start * ratio + (whisper_start - margin)
+            w_end = matched[-1].end * ratio + (whisper_start - margin)
+            new_words.append({"word": w["word"], "start": round(max(0, w_start), 3), "end": round(max(0, w_end), 3)})
+            span_idx += n_chars
+
+        return new_words if len(new_words) == len(words) else words
+    except Exception as e:
+        print(f"  [align] skipped due to: {type(e).__name__}: {e}")
+        return words
+
 
 @app.get("/health")
 async def health():
@@ -49,9 +114,6 @@ async def transcribe(websocket: WebSocket):
             tmp_path = tmp.name
         print(f"Step 3: saved to {tmp_path}")
 
-        await websocket.send_json({"status": "cleaning"})
-        print("Step 3.5: cleaning audio with DeepFilterNet")
-
         audio_to_transcribe = tmp_path
 
         if USE_DEEPFILTERNET:
@@ -76,32 +138,47 @@ async def transcribe(websocket: WebSocket):
             word_timestamps=True,
             vad_filter=True,
             vad_parameters=dict(
-                min_silence_duration_ms=500,
+                min_silence_duration_ms=300,
                 speech_pad_ms=200,
-                threshold=0.2,
+                threshold=0.3,
             ),
         )
 
         print("Step 5: transcription done")
 
-        await websocket.send_json({"status": "processing"})
+        await websocket.send_json({"status": "aligning"})
+        print("Step 5.5: loading audio for word-alignment")
+
+        align_waveform, sr = torchaudio.load(audio_to_transcribe)
+        if sr != 16000:
+            align_waveform = torchaudio.functional.resample(align_waveform, sr, 16000)
+        if align_waveform.shape[0] > 1:
+            align_waveform = align_waveform.mean(dim=0, keepdim=True)
 
         segments = []
         for seg in segments_raw:
             words = list(seg.words)
             if not words:
                 continue
+            word_dicts = [{"word": w.word, "start": w.start, "end": w.end} for w in words]
+
+            realigned = realign_word_timestamps(
+                align_waveform, word_dicts, word_dicts[0]["start"], word_dicts[-1]["end"]
+            )
+
             segments.append({
                 "text": seg.text.strip(),
-                "start": words[0].start,
-                "end": words[-1].end,
-                "words": [{"word": w.word, "start": w.start, "end": w.end} for w in words],
+                "start": realigned[0]["start"],
+                "end": realigned[-1]["end"],
+                "words": realigned,
             })
 
         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")
 
+        await websocket.send_json({"status": "processing"})
+
         srt_content = process(segments)
         print(f"Step 7: SRT built")