AliJabbarpour commited on
Commit
cfb7294
·
verified ·
1 Parent(s): 8a59102

Delete app.py

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