Update app.py

app.py

@@ -7,8 +7,10 @@ from transformers import AutoModelForCTC, Wav2Vec2BertProcessor
 from pytube import YouTube
 from transformers import pipeline
 import re
-import numpy as np
+from pydub import AudioSegment
+from scipy.io import wavfile
 from scipy.signal import wiener
+import numpy as np
 
 # pipe = pipeline(model="anzorq/w2v-bert-2.0-kbd", device=0) # old model
 pipe = pipeline(model="anzorq/w2v-bert-2.0-kbd-v2", device=0) # new model with a new tokenizer
@@ -26,19 +28,20 @@ reverse_pattern = re.compile('|'.join(re.escape(key) for key in reverse_replacem
 def replace_symbols_back(text):
     return reverse_pattern.sub(lambda match: reverse_replacements[match.group(0)], text)
 
-def normalize_audio(audio_tensor):
-    peak = torch.max(torch.abs(audio_tensor))
-    normalized_audio = audio_tensor / peak
-    return normalized_audio
+def normalize_audio(audio_path):
+    audio = AudioSegment.from_file(audio_path, format="mp4")
+    normalized_audio = audio.normalize()
+    normalized_audio.export(audio_path, format="mp4")
 
-def apply_wiener_filter(audio_tensor):
-    audio_data = audio_tensor.numpy()
+def apply_wiener_filter(audio_path):
+    sample_rate, audio_data = wavfile.read(audio_path)
     filtered_audio = wiener(audio_data)
-    return torch.tensor(filtered_audio)
+    wavfile.write(audio_path, sample_rate, filtered_audio.astype(np.int16))
 
-def resample_audio(audio_tensor, original_sample_rate, target_sample_rate=16000):
-    resampled_audio = torchaudio.transforms.Resample(original_sample_rate, target_sample_rate)(audio_tensor)
-    return resampled_audio
+def resample_audio(audio_path, target_sample_rate=16000):
+    audio, sample_rate = torchaudio.load(audio_path)
+    resampled_audio = torchaudio.transforms.Resample(sample_rate, target_sample_rate)(audio)
+    torchaudio.save(audio_path, resampled_audio, target_sample_rate)
 
 @spaces.GPU
 def transcribe_speech(audio, progress=gr.Progress()):
@@ -54,25 +57,20 @@ def transcribe_from_youtube(url, apply_improvements, progress=gr.Progress()):
     progress(0, "Downloading YouTube audio...")
     audio_path = YouTube(url).streams.filter(only_audio=True)[0].download(filename="tmp.mp4")
 
-    try:
-        audio, original_sample_rate = torchaudio.load(audio_path)
-
-        if apply_improvements:
-            progress(0.2, "Normalizing audio...")
-            audio = normalize_audio(audio)
+    if apply_improvements:
+        progress(0.2, "Normalizing audio...")
+        normalize_audio(audio_path)
 
-            progress(0.4, "Applying Wiener filter...")
-            audio = apply_wiener_filter(audio)
+        progress(0.4, "Applying Wiener filter...")
+        apply_wiener_filter(audio_path)
 
-            progress(0.6, "Resampling audio...")
-            audio = resample_audio(audio, original_sample_rate)
+        progress(0.6, "Resampling audio...")
+        resample_audio(audio_path)
 
-        progress(0.8, "Transcribing audio...")
-        transcription = transcribe_speech(audio)
+    progress(0.8, "Transcribing audio...")
+    transcription = transcribe_speech(audio_path)
 
-    finally:
-        if os.path.exists(audio_path):
-            os.remove(audio_path)
+    os.remove(audio_path)
 
     return transcription
 
@@ -117,4 +115,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         transcribe_button.click(fn=transcribe_from_youtube, inputs=[youtube_url, apply_improvements], outputs=transcription_output)
         youtube_url.change(populate_metadata, inputs=[youtube_url], outputs=[img, title])
 
-demo.launch()
+demo.launch()