| import os |
| import tempfile |
| import numpy as np |
| import torch |
| 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 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 |
|
|
| @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"}) |
|
|
| 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": "cleaning"}) |
| print("Step 3.5: cleaning audio with DeepFilterNet") |
|
|
| audio_to_transcribe = tmp_path |
|
|
| if USE_DEEPFILTERNET: |
| await websocket.send_json({"status": "cleaning"}) |
| 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)") |
|
|
| await websocket.send_json({"status": "transcribing"}) |
| print("Step 4: starting transcription") |
|
|
| segments_raw, info = model.transcribe( |
| audio_to_transcribe, |
| language="en", |
| word_timestamps=True, |
| vad_filter=True, |
| vad_parameters=dict( |
| min_silence_duration_ms=500, |
| speech_pad_ms=200, |
| threshold=0.1 |
| ), |
| ) |
|
|
| print("Step 5: transcription done") |
|
|
| await websocket.send_json({"status": "processing"}) |
|
|
| segments = [] |
| for seg in segments_raw: |
| words = list(seg.words) |
| if not words: |
| continue |
| 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], |
| }) |
|
|
| 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") |
|
|
| srt_content = process(segments) |
| print(f"Step 7: SRT built") |
|
|
| await websocket.send_json({ |
| "status": "done", |
| "srt": srt_content, |
| "language": info.language, |
| "duration": info.duration, |
| }) |
| print("Step 8: done sent!") |
|
|
| 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 |
