Upload app.py

app.py

@@ -0,0 +1,340 @@
+import os
+import asyncio
+import tempfile
+import torch
+import torchaudio
+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()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+print("Loading Whisper model...", flush=True)
+model = WhisperModel("large-v3", device="cpu", compute_type="int8")
+print("Whisper model ready!")
+
+print("Loading DeepFilterNet model...", flush=True)
+df_model, df_state, _ = init_df()
+print("DeepFilterNet model ready!")
+
+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!")
+
+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!")
+
+MIN_NOISE_BLOCK_SEC = 1.0  # حداقل طول یه بلوک بی‌گوینده تا «نویز» حساب بشه
+HEARTBEAT_SEC = 3  # هر چند ثانیه یه پیام زنده بفرستیم تا اتصال ساکت نمونه
+
+
+def get_speaker_at(diarization, time_sec):
+    for turn, _, speaker in diarization.itertracks(yield_label=True):
+        if turn.start <= time_sec <= turn.end:
+            return speaker
+    return None
+
+
+def find_leading_noise_cutoff(diarization, whisper_start, search_back_sec=4.0):
+    """
+    برمی‌گرده به عقب از شروعِ حدسیِ Whisper، دنبال یه بلوک نویز/بی‌گوینده‌ی
+    طولانی (حداقل MIN_NOISE_BLOCK_SEC) که دقیقاً قبل از این جمله باشه.
+    نوسان‌های خیلی کوتاه (کمتر از STABLE_SPEAKER_SEC) که یه گوینده برای یه
+    لحظه‌ی گذرا ظاهر میشه رو نادیده می‌گیره - فقط با دیدنِ یه گوینده‌ی
+    "پایدار" واقعاً متوقف میشه.
+    اگه پیدا کرد، زمان پایانِ اون بلوک نویز رو برمی‌گردونه (یعنی «از اینجا به بعد
+    اعتماد کن»)؛ اگه پیدا نکرد، None برمی‌گردونه (یعنی چیزی رو نبُر).
+    """
+    STABLE_SPEAKER_SEC = 0.2
+    window_start = max(0, whisper_start - search_back_sec)
+    step = 0.05
+    t = whisper_start
+    noise_end = None
+    noise_start = None
+    speaker_run_start = None
+
+    while t > window_start:
+        speaker = get_speaker_at(diarization, t)
+
+        if speaker is None:
+            speaker_run_start = None
+            if noise_end is None:
+                noise_end = t
+            noise_start = t
+        else:
+            if speaker_run_start is None:
+                speaker_run_start = t
+            run_len = speaker_run_start - t
+            if run_len >= STABLE_SPEAKER_SEC:
+                break
+            if noise_end is None:
+                noise_end = t
+            noise_start = t
+
+        t -= step
+
+    if noise_start is not None and noise_end is not None:
+        block_len = noise_end - noise_start
+        if block_len >= MIN_NOISE_BLOCK_SEC:
+            return noise_end
+    return None
+
+
+def realign_word_timestamps(waveform_16k, words, whisper_start, whisper_end, cutoff_start=None):
+    """
+    برای یه جمله، به‌جای زمان‌بندی حدسی خودِ Whisper، با یه مدل تخصصی جدا
+    (wav2vec2 از طریق torchaudio) دقیقاً چک می‌کنه هر کلمه از کجای صدا گفته شده.
+    اگه cutoff_start داده بشه (یعنی یه بلوک نویز قبل از جمله شناسایی شده)،
+    جستجو از همون نقطه شروع میشه، نه از قبل‌ترش.
+    اگه به هر دلیلی هم‌ترازی جواب نداد، همون زمان‌بندی قبلی Whisper دست‌نخورده می‌مونه.
+    """
+    try:
+        margin = 0.5
+        raw_start = whisper_start - margin
+        effective_start = max(raw_start, cutoff_start) if cutoff_start is not None else raw_start
+
+        start_sample = max(0, int(effective_start * 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"].lower() 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
+        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 + effective_start
+            w_end = matched[-1].end * ratio + effective_start
+            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
+
+
+def run_transcription(audio_path):
+    """فقط تشخیص گفتار - بلاک‌کننده، توی یه thread جدا اجرا میشه"""
+    return model.transcribe(
+        audio_path,
+        language="en",
+        word_timestamps=True,
+        vad_filter=True,
+        vad_parameters=dict(
+            min_silence_duration_ms=300,
+            speech_pad_ms=200,
+            threshold=0.3,
+        ),
+    )
+
+
+def run_diarization(audio_path):
+    """فقط تشخیص گوینده - بلاک‌کننده، توی یه thread جدا اجرا میشه"""
+    return diarize_pipeline(audio_path)
+
+
+def run_alignment_pass(segments_raw, align_waveform, diarization):
+    """
+    برای همه‌ی جمله‌ها، هم‌ترازی + تشخیص گوینده + برش نویز رو اجرا می‌کنه.
+    بلاک‌کننده‌ست (حلقه‌ی سنگین روی هر جمله)، توی یه thread جدا اجرا میشه.
+    """
+    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]
+
+        cutoff = find_leading_noise_cutoff(diarization, word_dicts[0]["start"])
+        if cutoff is not None:
+            print(f"  [noise-trim] '{seg.text.strip()[:40]}...' cutoff moved from {word_dicts[0]['start']:.2f}s to {cutoff:.2f}s")
+
+        realigned = realign_word_timestamps(
+            align_waveform, word_dicts, word_dicts[0]["start"], word_dicts[-1]["end"], cutoff_start=cutoff
+        )
+
+        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,
+        })
+    return segments
+
+
+async def run_with_heartbeat(websocket, blocking_fn, *args, status, detail):
+    """
+    یه تابع سنگین و بلاک‌کننده رو توی یه thread جدا اجرا می‌کنه، و تا وقتی تموم بشه،
+    هر چند ثانیه یه پیام "هنوز کار می‌کنم" روی همون status می‌فرسته - این هم جلوی
+    قطع‌شدن اتصال به‌خاطر سکوت طولانی رو می‌گیره، هم UI رو زنده نگه می‌داره.
+    """
+    task = asyncio.create_task(asyncio.to_thread(blocking_fn, *args))
+    elapsed = 0
+    await websocket.send_json({"status": status, "detail": detail})
+    while not task.done():
+        try:
+            await asyncio.wait_for(asyncio.shield(task), timeout=HEARTBEAT_SEC)
+        except asyncio.TimeoutError:
+            elapsed += HEARTBEAT_SEC
+            await websocket.send_json({
+                "status": status,
+                "detail": f"{detail} ({elapsed} ثانیه)",
+            })
+    return await task
+
+
+@app.get("/health")
+async def health():
+    return {"status": "ok", "message": "SubFlow backend is running"}
+
+
+@app.websocket("/transcribe")
+async def transcribe(websocket: WebSocket):
+    await websocket.accept()
+    tmp_path = None
+    enhanced_path = None
+
+    try:
+        print("Step 1: waiting for bytes")
+        audio_bytes = await websocket.receive_bytes()
+        print(f"Step 2: received {len(audio_bytes)} bytes")
+        await websocket.send_json({"status": "receiving", "detail": "فایل دریافت شد"})
+
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+            tmp.write(audio_bytes)
+            tmp_path = tmp.name
+        print(f"Step 3: saved to {tmp_path}")
+
+        audio_to_transcribe = tmp_path
+
+        if USE_DEEPFILTERNET:
+            await websocket.send_json({"status": "cleaning", "detail": "در حال حذف نویز پس‌زمینه"})
+            print("Step 3.5: cleaning audio with DeepFilterNet")
+
+            audio_data, _ = load_audio(tmp_path, sr=df_state.sr())
+            enhanced_audio = enhance(df_model, df_state, audio_data, atten_lim_db=20)
+            enhanced_path = tmp_path.replace(".wav", "_enhanced.wav")
+            save_audio(enhanced_path, enhanced_audio, df_state.sr())
+            audio_to_transcribe = enhanced_path
+            print(f"Step 3.6: cleaned audio saved to {enhanced_path}")
+        else:
+            print("Step 3.5: skipping DeepFilterNet (disabled for A/B test)")
+
+        print("Step 4: starting transcription")
+        segments_raw, info = await run_with_heartbeat(
+            websocket, run_transcription, audio_to_transcribe,
+            status="transcribing", detail="مدل large-v3 در حال شنیدن صداست",
+        )
+        print("Step 5: transcription done")
+
+        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)
+
+        print("Step 5.5: running speaker diarization")
+        diarization = await run_with_heartbeat(
+            websocket, run_diarization, audio_to_transcribe,
+            status="aligning", detail="در حال تشخیص گوینده‌ها",
+        )
+        print("Step 5.6: diarization done, now aligning words per segment")
+
+        segments_raw_list = list(segments_raw)
+        segments = await run_with_heartbeat(
+            websocket, run_alignment_pass, segments_raw_list, align_waveform, diarization,
+            status="aligning", detail=f"در حال هماهنگ‌سازی زمان‌بندی {len(segments_raw_list)} جمله",
+        )
+
+        print(f"Step 6: got {len(segments)} segments")
+        for i, seg in enumerate(segments):
+            print(f"  seg[{i}]: '{seg['text']}' | speaker={seg['speaker']} | dur={seg['end']-seg['start']:.2f}s")
+
+        await websocket.send_json({
+            "status": "processing",
+            "detail": f"در حال چیدمان {len(segments)} خط زیرنویس",
+        })
+
+        srt_content = process(segments)
+        print(f"Step 7: SRT built")
+
+        await websocket.send_json({
+            "status": "done",
+            "detail": "زیرنویس آماده شد",
+            "srt": srt_content,
+            "language": info.language,
+            "duration": info.duration,
+        })
+        print("Step 8: done sent!")
+        await asyncio.sleep(1)
+
+    except WebSocketDisconnect:
+        print("WebSocket disconnected by client")
+    except Exception as e:
+        print(f"ERROR: {type(e).__name__}: {e}")
+        import traceback
+        traceback.print_exc()
+        try:
+            await websocket.send_json({"status": "error", "message": str(e)})
+        except Exception:
+            pass
+    finally:
+        if tmp_path and os.path.exists(tmp_path):
+            os.remove(tmp_path)
+        if enhanced_path and os.path.exists(enhanced_path):
+            os.remove(enhanced_path)
+        try:
+            await websocket.close()
+        except Exception:
+            pass