Update app.py

app.py

@@ -1,5 +1,7 @@
 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
@@ -15,7 +17,6 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-import sys
 print("Loading Whisper model...", flush=True)
 model = WhisperModel("large-v3", device="cpu", compute_type="int8")
 print("Whisper model ready!")
@@ -24,10 +25,39 @@ print("Loading DeepFilterNet model...", flush=True)
 df_model, df_state, _ = init_df()
 print("DeepFilterNet model ready!")
 
+
+def compress_dynamic_range(audio, sr, threshold_db=-28, ratio=4.0, makeup_gain_db=8, block_ms=20):
+    """
+    صدای آروم رو نسبت به صدای بلند بلندتر می‌کنه (یکسان‌سازی حجم صدا)،
+    تا صدای یه نفر که از میکروفون دوره به صدای کسی که میکروفون دستشه نزدیک‌تر بشه.
+    """
+    audio = audio.astype(np.float32)
+    block_size = max(1, int(sr * block_ms / 1000))
+    n_blocks = int(np.ceil(len(audio) / block_size))
+    pad_len = n_blocks * block_size - len(audio)
+    padded = np.pad(audio, (0, pad_len))
+    blocks = padded.reshape(n_blocks, block_size)
+
+    rms = np.sqrt(np.mean(blocks ** 2, axis=1) + 1e-10)
+    rms_db = 20 * np.log10(rms + 1e-10)
+
+    target_db = np.where(rms_db > threshold_db, threshold_db + (rms_db - threshold_db) / ratio, rms_db)
+    gain_db = target_db - rms_db + makeup_gain_db
+    gain = 10 ** (gain_db / 20)
+
+    kernel = np.ones(5, dtype=np.float32) / 5
+    gain_smooth = np.convolve(gain, kernel, mode="same")
+
+    gain_per_sample = np.repeat(gain_smooth, block_size)[:len(audio)].astype(np.float32)
+    compressed = audio * gain_per_sample
+    return np.clip(compressed, -1.0, 1.0).astype(np.float32)
+
+
 @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()
@@ -46,9 +76,12 @@ async def transcribe(websocket: WebSocket):
         print(f"Step 3: saved to {tmp_path}")
 
         await websocket.send_json({"status": "cleaning"})
-        print("Step 3.5: cleaning audio with DeepFilterNet")
+        print("Step 3.5: leveling volume + cleaning audio with DeepFilterNet")
 
         audio_data, _ = load_audio(tmp_path, sr=df_state.sr())
+        audio_np = compress_dynamic_range(audio_data.numpy()[0], df_state.sr())
+        audio_data = torch.from_numpy(audio_np).unsqueeze(0)
+
         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())
@@ -67,7 +100,7 @@ async def transcribe(websocket: WebSocket):
                 speech_pad_ms=200,
             ),
         )
-        
+
         print("Step 5: transcription done")
 
         await websocket.send_json({"status": "processing"})
@@ -117,5 +150,4 @@ async def transcribe(websocket: WebSocket):
         try:
             await websocket.close()
         except Exception:
-            pass
-            
+            pass