AliJabbarpour commited on
Commit
bcf2b95
·
verified ·
1 Parent(s): b3278a7

Delete app.py

Browse files
Files changed (1) hide show
  1. app.py +0 -400
app.py DELETED
@@ -1,400 +0,0 @@
1
- import os
2
- import asyncio
3
- import tempfile
4
- import torch
5
- import torchaudio
6
-
7
- torch.set_num_threads(2) # جلوی رقابت بی‌فایده‌ی رشته‌ها روی CPU محدود رو می‌گیره
8
- from fastapi import FastAPI, WebSocket, WebSocketDisconnect
9
- from fastapi.middleware.cors import CORSMiddleware
10
- from faster_whisper import WhisperModel
11
- from df.enhance import enhance, init_df, load_audio, save_audio
12
- from pyannote.audio import Pipeline as DiarizationPipeline
13
- from pipeline.postprocess import process
14
-
15
- app = FastAPI()
16
-
17
- app.add_middleware(
18
- CORSMiddleware,
19
- allow_origins=["*"],
20
- allow_methods=["*"],
21
- allow_headers=["*"],
22
- )
23
-
24
- print("Loading Whisper model...", flush=True)
25
- model = WhisperModel("large-v3", device="cpu", compute_type="int8")
26
- print("Whisper model ready!")
27
-
28
- print("Loading DeepFilterNet model...", flush=True)
29
- df_model, df_state, _ = init_df()
30
- print("DeepFilterNet model ready!")
31
-
32
- USE_DEEPFILTERNET = False # غیرفعال - شواهد نشون داد کمکی نمی‌کنه و بعضی صداها رو خراب می‌کنه
33
-
34
- print("Loading word-alignment model...", flush=True)
35
- align_bundle = torchaudio.pipelines.MMS_FA
36
- align_model = align_bundle.get_model().to("cpu")
37
- align_labels = align_bundle.get_labels()
38
- align_dictionary = {c: i for i, c in enumerate(align_labels)}
39
- print("Word-alignment model ready!")
40
-
41
- print("Loading speaker-diarization model...", flush=True)
42
- HF_TOKEN = os.environ.get("HF_TOKEN")
43
- diarize_pipeline = DiarizationPipeline.from_pretrained(
44
- "pyannote/speaker-diarization-3.1", use_auth_token=HF_TOKEN
45
- )
46
- print("Speaker-diarization model ready!")
47
-
48
- MIN_NOISE_BLOCK_SEC = 1.0 # حداقل طول یه بلوک بی‌گوینده تا «نویز» حساب بشه
49
- HEARTBEAT_SEC = 3 # هر چند ثانیه یه پیام زنده بفرستیم تا اتصال ساکت نمونه
50
-
51
-
52
- def get_speaker_at(diarization, time_sec):
53
- for turn, _, speaker in diarization.itertracks(yield_label=True):
54
- if turn.start <= time_sec <= turn.end:
55
- return speaker
56
- return None
57
-
58
-
59
- def get_speaker_for_segment(diarization, seg_start, seg_end):
60
- """
61
- به‌جای اعتماد به یه لحظه‌ی تنها (که ممکنه نزدیک مرز بین دو گوینده باشه)،
62
- چند نقطه در طول جمله رو چک می‌کنه و گوینده‌ای که بیشترین رأی رو آورده
63
- برمی‌گردونه - این دقت رو بهتر می‌کنه ولی صد در صد تضمین‌شده نیست.
64
- """
65
- duration = seg_end - seg_start
66
- if duration <= 0:
67
- return get_speaker_at(diarization, seg_start)
68
-
69
- sample_points = [seg_start + duration * frac for frac in (0.2, 0.4, 0.5, 0.6, 0.8)]
70
- votes = {}
71
- for t in sample_points:
72
- speaker = get_speaker_at(diarization, t)
73
- if speaker is not None:
74
- votes[speaker] = votes.get(speaker, 0) + 1
75
-
76
- if not votes:
77
- return None
78
- return max(votes, key=votes.get)
79
-
80
-
81
- def find_leading_noise_cutoff(diarization, whisper_start, search_back_sec=4.0):
82
- """
83
- برمی‌گرده به عقب از شروعِ حدسیِ Whisper، دنبال یه بلوک نویز/بی‌گوینده‌ی
84
- طولانی (حداقل MIN_NOISE_BLOCK_SEC) که دقیقاً قبل از این جمله باشه.
85
- نوسان‌های خیلی کوتاه (کمتر از STABLE_SPEAKER_SEC) که یه گوینده برای یه
86
- لحظه‌ی گذرا ظاهر میشه رو نادیده می‌گیره - فقط با دیدنِ یه گوینده‌ی
87
- "پایدار" واقعاً متوقف میشه.
88
- اگه پیدا کرد، زمان پایانِ اون بلوک نویز رو برمی‌گردونه (یعنی «از اینجا به بعد
89
- اعتماد کن»)؛ اگه پیدا نکرد، None برمی‌گردونه (یعنی چیزی رو نبُر).
90
- """
91
- STABLE_SPEAKER_SEC = 0.2
92
- window_start = max(0, whisper_start - search_back_sec)
93
- step = 0.05
94
- t = whisper_start
95
- noise_end = None
96
- noise_start = None
97
- speaker_run_start = None
98
-
99
- while t > window_start:
100
- speaker = get_speaker_at(diarization, t)
101
-
102
- if speaker is None:
103
- speaker_run_start = None
104
- if noise_end is None:
105
- noise_end = t
106
- noise_start = t
107
- else:
108
- if speaker_run_start is None:
109
- speaker_run_start = t
110
- run_len = speaker_run_start - t
111
- if run_len >= STABLE_SPEAKER_SEC:
112
- break
113
- if noise_end is None:
114
- noise_end = t
115
- noise_start = t
116
-
117
- t -= step
118
-
119
- if noise_start is not None and noise_end is not None:
120
- block_len = noise_end - noise_start
121
- if block_len >= MIN_NOISE_BLOCK_SEC:
122
- return noise_end
123
- return None
124
-
125
-
126
- def realign_word_timestamps(waveform_16k, words, whisper_start, whisper_end, cutoff_start=None):
127
- """
128
- برای یه جمله، به‌جای ��مان‌بندی حدسی خودِ Whisper، با یه مدل تخصصی جدا
129
- (wav2vec2 از طریق torchaudio) دقیقاً چک می‌کنه هر کلمه از کجای صدا گفته شده.
130
- اگه cutoff_start داده بشه (یعنی یه بلوک نویز قبل از جمله شناسایی شده)،
131
- جستجو از همون نقطه شروع میشه، نه از قبل‌ترش.
132
- اگه به هر دلیلی هم‌ترازی جواب نداد، همون زمان‌بندی قبلی Whisper دست‌نخورده می‌مونه.
133
- """
134
- try:
135
- margin = 0.5
136
- raw_start = whisper_start - margin
137
- effective_start = max(raw_start, cutoff_start) if cutoff_start is not None else raw_start
138
-
139
- start_sample = max(0, int(effective_start * 16000))
140
- end_sample = int((whisper_end + margin) * 16000)
141
- clip = waveform_16k[:, start_sample:end_sample]
142
-
143
- clean_words = []
144
- for w in words:
145
- token = "".join(ch for ch in w["word"].lower() if ch in align_dictionary)
146
- if token:
147
- clean_words.append((w, token))
148
-
149
- if not clean_words:
150
- return words
151
-
152
- transcript = "".join(tok for _, tok in clean_words)
153
- tokens = [align_dictionary[c] for c in transcript]
154
-
155
- with torch.inference_mode():
156
- emission, _ = align_model(clip)
157
-
158
- targets = torch.tensor([tokens], dtype=torch.int32)
159
- input_lengths = torch.tensor([emission.shape[1]])
160
- target_lengths = torch.tensor([len(tokens)])
161
-
162
- aligned_tokens, scores = torchaudio.functional.forced_align(
163
- emission, targets, input_lengths, target_lengths, blank=0
164
- )
165
- spans = torchaudio.functional.merge_tokens(aligned_tokens[0], scores[0])
166
-
167
- ratio = clip.shape[1] / emission.shape[1] / 16000
168
- new_words = []
169
- span_idx = 0
170
- for w, tok in clean_words:
171
- n_chars = len(tok)
172
- matched = [s for s in spans[span_idx:] if s.token != 0][:n_chars]
173
- if len(matched) < n_chars:
174
- new_words.append(w)
175
- continue
176
- w_start = matched[0].start * ratio + effective_start
177
- w_end = matched[-1].end * ratio + effective_start
178
- new_words.append({"word": w["word"], "start": round(max(0, w_start), 3), "end": round(max(0, w_end), 3)})
179
- span_idx += n_chars
180
-
181
- return new_words if len(new_words) == len(words) else words
182
- except Exception as e:
183
- print(f" [align] skipped due to: {type(e).__name__}: {e}")
184
- return words
185
-
186
-
187
- def run_transcription(audio_path):
188
- """فقط تشخیص گفتار - بلاک‌کننده، توی یه thread جدا اجرا میشه"""
189
- return model.transcribe(
190
- audio_path,
191
- language="en",
192
- word_timestamps=True,
193
- vad_filter=True,
194
- vad_parameters=dict(
195
- min_silence_duration_ms=300,
196
- speech_pad_ms=200,
197
- threshold=0.3,
198
- ),
199
- )
200
-
201
-
202
- def run_diarization(audio_path):
203
- """فقط تشخیص گوینده - بلاک‌کننده، توی یه thread جدا اجرا میشه"""
204
- return diarize_pipeline(audio_path, max_speakers=8)
205
-
206
-
207
- def run_alignment_pass(segments_raw, align_waveform, diarization):
208
- """
209
- برای همه‌ی جمله‌ها، هم‌ترازی + تشخیص گوینده + برش نویز رو اجرا می‌کنه.
210
- بلاک‌کننده‌ست (حلقه‌ی سنگین روی هر جمله)، توی یه thread جدا اجرا میشه.
211
- """
212
- segments = []
213
- for seg in segments_raw:
214
- words = list(seg.words)
215
- if not words:
216
- continue
217
- word_dicts = [{"word": w.word, "start": w.start, "end": w.end} for w in words]
218
-
219
- cutoff = find_leading_noise_cutoff(diarization, word_dicts[0]["start"])
220
- if cutoff is not None:
221
- print(f" [noise-trim] '{seg.text.strip()[:40]}...' cutoff moved from {word_dicts[0]['start']:.2f}s to {cutoff:.2f}s")
222
-
223
- realigned = realign_word_timestamps(
224
- align_waveform, word_dicts, word_dicts[0]["start"], word_dicts[-1]["end"], cutoff_start=cutoff
225
- )
226
-
227
- speaker = get_speaker_for_segment(diarization, realigned[0]["start"], realigned[-1]["end"])
228
-
229
- segments.append({
230
- "text": seg.text.strip(),
231
- "start": realigned[0]["start"],
232
- "end": realigned[-1]["end"],
233
- "words": realigned,
234
- "speaker": speaker,
235
- })
236
- return segments
237
-
238
-
239
- class ProcessingTimeout(Exception):
240
- pass
241
-
242
-
243
- async def run_with_heartbeat(websocket, blocking_fn, *args, status, detail, max_seconds=240):
244
- """
245
- یه تابع سنگین و بلاک‌کننده رو توی یه thread جدا اجرا می‌کنه، و تا وقتی تموم بشه،
246
- هر چند ثانیه یه پیام "هنوز کار می‌کنم" روی همون status می‌فرسته - این هم جلوی
247
- قطع‌شدن اتصال به‌خاطر سکوت طولانی رو می‌گیره، هم UI رو زنده نگه می‌داره.
248
- اگه از max_seconds بیشتر طول بکشه، به‌جای قفل ابدی، خطای مشخص میده.
249
- """
250
- task = asyncio.create_task(asyncio.to_thread(blocking_fn, *args))
251
- elapsed = 0
252
- await websocket.send_json({"status": status, "detail": detail})
253
- while not task.done():
254
- try:
255
- await asyncio.wait_for(asyncio.shield(task), timeout=HEARTBEAT_SEC)
256
- except asyncio.TimeoutError:
257
- elapsed += HEARTBEAT_SEC
258
- if elapsed >= max_seconds:
259
- task.cancel()
260
- try:
261
- await task
262
- except (asyncio.CancelledError, Exception):
263
- pass
264
- raise ProcessingTimeout(
265
- f"مرحله‌ی «{detail}» بیشتر از {max_seconds} ثانیه طول کشید و متوقف شد."
266
- )
267
- await websocket.send_json({
268
- "status": status,
269
- "detail": f"{detail} ({elapsed} ثانیه)",
270
- })
271
- return await task
272
-
273
-
274
- @app.get("/health")
275
- async def health():
276
- return {"status": "ok", "message": "SubFlow backend is running"}
277
-
278
-
279
- @app.websocket("/transcribe")
280
- async def transcribe(websocket: WebSocket):
281
- await websocket.accept()
282
- tmp_path = None
283
- enhanced_path = None
284
-
285
- try:
286
- await websocket.send_json({"status": "receiving", "detail": "در انتظار دریافت فایل..."})
287
- print("Step 1: waiting for chunks")
288
- chunks = []
289
- total_bytes = 0
290
- while True:
291
- message = await websocket.receive()
292
- if "bytes" in message and message["bytes"] is not None:
293
- chunk = message["bytes"]
294
- chunks.append(chunk)
295
- total_bytes += len(chunk)
296
- elif "text" in message and message["text"] == "__END__":
297
- break
298
- audio_bytes = b"".join(chunks)
299
- print(f"Step 2: received {total_bytes} bytes in {len(chunks)} chunks")
300
- size_mb = len(audio_bytes) / (1024 * 1024)
301
- await websocket.send_json({"status": "receiving", "detail": f"{size_mb:.1f} مگابایت دریافت شد"})
302
-
303
- with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
304
- tmp.write(audio_bytes)
305
- tmp_path = tmp.name
306
- print(f"Step 3: saved to {tmp_path}")
307
- await websocket.send_json({"status": "receiving", "detail": "فایل آماده‌ی پردازش شد"})
308
-
309
- audio_to_transcribe = tmp_path
310
-
311
- if USE_DEEPFILTERNET:
312
- await websocket.send_json({"status": "cleaning", "detail": "در حال حذف نویز پس‌زمینه"})
313
- print("Step 3.5: cleaning audio with DeepFilterNet")
314
-
315
- audio_data, _ = load_audio(tmp_path, sr=df_state.sr())
316
- enhanced_audio = enhance(df_model, df_state, audio_data, atten_lim_db=20)
317
- enhanced_path = tmp_path.replace(".wav", "_enhanced.wav")
318
- save_audio(enhanced_path, enhanced_audio, df_state.sr())
319
- audio_to_transcribe = enhanced_path
320
- print(f"Step 3.6: cleaned audio saved to {enhanced_path}")
321
- else:
322
- print("Step 3.5: skipping DeepFilterNet (disabled for A/B test)")
323
-
324
- print("Step 4: starting transcription")
325
- segments_raw, info = await run_with_heartbeat(
326
- websocket, run_transcription, audio_to_transcribe,
327
- status="transcribing", detail="مدل large-v3 در حال شنیدن صداست",
328
- )
329
- print("Step 5: transcription done")
330
- await websocket.send_json({"status": "transcribing", "detail": "تشخیص گفتار کامل شد"})
331
-
332
- align_waveform, sr = torchaudio.load(audio_to_transcribe)
333
- if sr != 16000:
334
- align_waveform = torchaudio.functional.resample(align_waveform, sr, 16000)
335
- if align_waveform.shape[0] > 1:
336
- align_waveform = align_waveform.mean(dim=0, keepdim=True)
337
-
338
- print("Step 5.5: running speaker diarization")
339
- diarization = await run_with_heartbeat(
340
- websocket, run_diarization, audio_to_transcribe,
341
- status="aligning", detail="در حال تشخیص گوینده‌ها",
342
- )
343
- print("Step 5.6: diarization done, now aligning words per segment")
344
- await websocket.send_json({"status": "aligning", "detail": "گوینده‌ها شناسایی شدند"})
345
-
346
- segments_raw_list = list(segments_raw)
347
- segments = await run_with_heartbeat(
348
- websocket, run_alignment_pass, segments_raw_list, align_waveform, diarization,
349
- status="aligning", detail=f"در حال هماهنگ‌سازی زمان‌بندی {len(segments_raw_list)} جمله",
350
- )
351
-
352
- print(f"Step 6: got {len(segments)} segments")
353
- await websocket.send_json({"status": "aligning", "detail": f"{len(segments)} خط شناسایی شد"})
354
- for i, seg in enumerate(segments):
355
- print(f" seg[{i}]: '{seg['text']}' | speaker={seg['speaker']} | dur={seg['end']-seg['start']:.2f}s")
356
-
357
- await websocket.send_json({
358
- "status": "processing",
359
- "detail": f"در حال چیدمان {len(segments)} خط زیرنویس",
360
- })
361
-
362
- srt_content = process(segments)
363
- print(f"Step 7: SRT built")
364
- await websocket.send_json({"status": "processing", "detail": "زیرنویس ساخته شد"})
365
-
366
- await websocket.send_json({
367
- "status": "done",
368
- "detail": "زیرنویس آماده شد",
369
- "srt": srt_content,
370
- "language": info.language,
371
- "duration": info.duration,
372
- })
373
- print("Step 8: done sent!")
374
- await asyncio.sleep(1)
375
-
376
- except WebSocketDisconnect:
377
- print("WebSocket disconnected by client")
378
- except ProcessingTimeout as e:
379
- print(f"TIMEOUT: {e}")
380
- try:
381
- await websocket.send_json({"status": "error", "message": str(e)})
382
- except Exception:
383
- pass
384
- except Exception as e:
385
- print(f"ERROR: {type(e).__name__}: {e}")
386
- import traceback
387
- traceback.print_exc()
388
- try:
389
- await websocket.send_json({"status": "error", "message": str(e)})
390
- except Exception:
391
- pass
392
- finally:
393
- if tmp_path and os.path.exists(tmp_path):
394
- os.remove(tmp_path)
395
- if enhanced_path and os.path.exists(enhanced_path):
396
- os.remove(enhanced_path)
397
- try:
398
- await websocket.close()
399
- except Exception:
400
- pass