Update app.py

app.py

@@ -36,7 +36,6 @@ def preprocess_audio(audio_tensor, original_sample_rate, apply_normalization):
 
     if apply_normalization:
         audio_tensor = audio_tensor / torch.max(torch.abs(audio_tensor)) # Normalize
-    # audio_tensor = torch.clamp(audio_tensor, min=-1, max=1)
 
     audio_tensor = torchaudio.functional.resample(audio_tensor, orig_freq=original_sample_rate, new_freq=16000)  # Resample
     return audio_tensor
@@ -52,13 +51,31 @@ def wiener_filter(audio_tensor):
     return torch.tensor(filtered_audio, dtype=audio_tensor.dtype)
 
 @spaces.GPU
-def transcribe_speech(audio, progress=gr.Progress()):
+def transcribe_speech(audio, apply_wiener_filter=False, apply_normalization=False, apply_spectral_gating=False, progress=gr.Progress()):
     if audio is None:
         return "No audio received.", None
-    progress(0.5, desc="Transcribing audio...")
-    audio_np = audio.numpy().squeeze()
+
+    progress(0.1, desc="Preprocessing audio...")
+    audio_tensor, original_sample_rate = torchaudio.load(audio)
+    audio_tensor = preprocess_audio(audio_tensor, original_sample_rate, apply_normalization)
+
+    if apply_wiener_filter:
+        progress(0.3, desc="Applying Wiener filter...")
+        audio_tensor = wiener_filter(audio_tensor)
+
+    if apply_spectral_gating:
+        progress(0.5, desc="Applying Spectral Gating filter...")
+        audio_tensor = spectral_gating(audio_tensor)
+
+    progress(0.7, desc="Transcribing audio...")
+    audio_np = audio_tensor.numpy().squeeze()
     transcription = pipe(audio_np, chunk_length_s=10)['text']
-    return replace_symbols_back(transcription), audio
+    transcription = replace_symbols_back(transcription)
+
+    audio_np = audio_tensor.numpy().squeeze()
+    sf.write("temp_audio.wav", audio_np, 16000, subtype='PCM_16')
+
+    return transcription, "temp_audio.wav"
 
 def transcribe_from_youtube(url, apply_wiener_filter, apply_normalization, apply_spectral_gating, progress=gr.Progress()):
     progress(0, "Downloading YouTube audio...")
@@ -70,20 +87,20 @@ def transcribe_from_youtube(url, apply_wiener_filter, apply_normalization, apply
         stream.stream_to_buffer(audio_data)
         audio_data.seek(0)
     
-        audio, original_sample_rate = torchaudio.load(audio_data)
-        audio = preprocess_audio(audio, original_sample_rate, apply_normalization)
+        audio_tensor, original_sample_rate = torchaudio.load(audio_data)
+        audio_tensor = preprocess_audio(audio_tensor, original_sample_rate, apply_normalization)
 
         if apply_wiener_filter:
             progress(0.4, "Applying Wiener filter...")
-            audio = wiener_filter(audio)
+            audio_tensor = wiener_filter(audio_tensor)
 
         if apply_spectral_gating:
             progress(0.4, "Applying Spectral Gating filter...")
-            audio = spectral_gating(audio)
+            audio_tensor = spectral_gating(audio_tensor)
 
-        transcription, _ = transcribe_speech(audio)
+        transcription, _ = transcribe_speech(audio_tensor)
 
-        audio_np = audio.numpy().squeeze()
+        audio_np = audio_tensor.numpy().squeeze()
         sf.write("temp_audio.wav", audio_np, 16000, subtype='PCM_16')
 
     except Exception as e:
@@ -116,7 +133,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         transcription_output = gr.Textbox(label="Transcription")
         audio_output = gr.Audio(label="Processed Audio")
 
-        transcribe_button.click(fn=transcribe_speech, inputs=mic_audio, outputs=[transcription_output, audio_output])
+        transcribe_button.click(fn=transcribe_speech, inputs=[mic_audio], outputs=[transcription_output, audio_output])
 
     with gr.Tab("YouTube URL"):
         gr.Markdown("## Transcribe speech from YouTube video")