Update app.py

app.py

@@ -7,10 +7,8 @@ from transformers import AutoModelForCTC, Wav2Vec2BertProcessor
 from pytube import YouTube
 from transformers import pipeline
 import re
-from pydub import AudioSegment
-from scipy.io import wavfile
-from scipy.signal import wiener
 import numpy as np
+from scipy.signal import wiener
 
 # 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
@@ -28,20 +26,19 @@ 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_path):
-    audio = AudioSegment.from_file(audio_path, format="mp4")
-    normalized_audio = audio.normalize()
-    normalized_audio.export(audio_path, format="mp4")
+def normalize_audio(audio_tensor):
+    peak = torch.max(torch.abs(audio_tensor))
+    normalized_audio = audio_tensor / peak
+    return normalized_audio
 
-def apply_wiener_filter(audio_path):
-    sample_rate, audio_data = wavfile.read(audio_path)
+def apply_wiener_filter(audio_tensor):
+    audio_data = audio_tensor.numpy()
     filtered_audio = wiener(audio_data)
-    wavfile.write(audio_path, sample_rate, filtered_audio.astype(np.int16))
+    return torch.tensor(filtered_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)
+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
 
 @spaces.GPU
 def transcribe_speech(audio, progress=gr.Progress()):
@@ -57,20 +54,25 @@ 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")
 
-    if apply_improvements:
-        progress(0.2, "Normalizing audio...")
-        normalize_audio(audio_path)
+    try:
+        audio, original_sample_rate = torchaudio.load(audio_path)
+
+        if apply_improvements:
+            progress(0.2, "Normalizing audio...")
+            audio = normalize_audio(audio)
 
-        progress(0.4, "Applying Wiener filter...")
-        apply_wiener_filter(audio_path)
+            progress(0.4, "Applying Wiener filter...")
+            audio = apply_wiener_filter(audio)
 
-        progress(0.6, "Resampling audio...")
-        resample_audio(audio_path)
+            progress(0.6, "Resampling audio...")
+            audio = resample_audio(audio, original_sample_rate)
 
-    progress(0.8, "Transcribing audio...")
-    transcription = transcribe_speech(audio_path)
+        progress(0.8, "Transcribing audio...")
+        transcription = transcribe_speech(audio)
 
-    os.remove(audio_path)
+    finally:
+        if os.path.exists(audio_path):
+            os.remove(audio_path)
 
     return transcription
 
@@ -115,4 +117,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()