| import os |
| import asyncio |
| import tempfile |
| import torch |
| import torchaudio |
|
|
| torch.set_num_threads(2) |
| 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 get_speaker_for_segment(diarization, seg_start, seg_end): |
| """ |
| به‌جای اعتماد به یه لحظه‌ی تنها (که ممکنه نزدیک مرز بین دو گوینده باشه)، |
| چند نقطه در طول جمله رو چک می‌کنه و گوینده‌ای که بیشترین رأی رو آورده |
| برمی‌گردونه - این دقت رو بهتر می‌کنه ولی صد در صد تضمین‌شده نیست. |
| """ |
| duration = seg_end - seg_start |
| if duration <= 0: |
| return get_speaker_at(diarization, seg_start) |
|
|
| sample_points = [seg_start + duration * frac for frac in (0.2, 0.4, 0.5, 0.6, 0.8)] |
| votes = {} |
| for t in sample_points: |
| speaker = get_speaker_at(diarization, t) |
| if speaker is not None: |
| votes[speaker] = votes.get(speaker, 0) + 1 |
|
|
| if not votes: |
| return None |
| return max(votes, key=votes.get) |
|
|
|
|
| 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, max_speakers=8) |
|
|
|
|
| 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 |
| ) |
|
|
| speaker = get_speaker_for_segment(diarization, realigned[0]["start"], realigned[-1]["end"]) |
|
|
| segments.append({ |
| "text": seg.text.strip(), |
| "start": realigned[0]["start"], |
| "end": realigned[-1]["end"], |
| "words": realigned, |
| "speaker": speaker, |
| }) |
| return segments |
|
|
|
|
| class ProcessingTimeout(Exception): |
| pass |
|
|
|
|
| async def run_with_heartbeat(websocket, blocking_fn, *args, status, detail, max_seconds=240): |
| """ |
| یه تابع سنگین و بلاک‌کننده رو توی یه thread جدا اجرا می‌کنه، و تا وقتی تموم بشه، |
| هر چند ثانیه یه پیام "هنوز کار می‌کنم" روی همون status می‌فرسته - این هم جلوی |
| قطع‌شدن اتصال به‌خاطر سکوت طولانی رو می‌گیره، هم UI رو زنده نگه می‌داره. |
| اگه از max_seconds بیشتر طول بکشه، به‌جای قفل ابدی، خطای مشخص میده. |
| """ |
| 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 |
| if elapsed >= max_seconds: |
| task.cancel() |
| raise ProcessingTimeout( |
| f"مرحله‌ی «{detail}» بیشتر از {max_seconds} ثانیه طول کشید و متوقف شد." |
| ) |
| 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: |
| await websocket.send_json({"status": "receiving", "detail": "در انتظار دریافت فایل..."}) |
| print("Step 1: waiting for chunks") |
| chunks = [] |
| total_bytes = 0 |
| while True: |
| message = await websocket.receive() |
| if "bytes" in message and message["bytes"] is not None: |
| chunk = message["bytes"] |
| chunks.append(chunk) |
| total_bytes += len(chunk) |
| elif "text" in message and message["text"] == "__END__": |
| break |
| audio_bytes = b"".join(chunks) |
| print(f"Step 2: received {total_bytes} bytes in {len(chunks)} chunks") |
| size_mb = len(audio_bytes) / (1024 * 1024) |
| await websocket.send_json({"status": "receiving", "detail": f"{size_mb:.1f} مگابایت دریافت شد"}) |
|
|
| 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}") |
| await websocket.send_json({"status": "receiving", "detail": "فایل آماده‌ی پردازش شد"}) |
|
|
| 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") |
| await websocket.send_json({"status": "transcribing", "detail": "تشخیص گفتار کامل شد"}) |
|
|
| 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") |
| await websocket.send_json({"status": "aligning", "detail": "گوینده‌ها شناسایی شدند"}) |
|
|
| 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") |
| await websocket.send_json({"status": "aligning", "detail": f"{len(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": "processing", "detail": "زیرنویس ساخته شد"}) |
|
|
| 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 ProcessingTimeout as e: |
| print(f"TIMEOUT: {e}") |
| try: |
| await websocket.send_json({"status": "error", "message": str(e)}) |
| except Exception: |
| pass |
| 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 |
|
|
