Update app.py

app.py

@@ -3,18 +3,16 @@ import os
 import gradio as gr
 import torch
 import torchaudio
-from transformers import AutoModelForCTC, Wav2Vec2BertProcessor
-from pytube import YouTube
 from transformers import pipeline
+from pytube import YouTube
 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
+from io import BytesIO
 
-# 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
+pipe = pipeline(model="anzorq/w2v-bert-2.0-kbd-v2", device=0)
 
+# Define the replacements for Kabardian transcription
 replacements = [
     ('гъ', 'ɣ'), ('дж', 'j'), ('дз', 'ӡ'), ('жь', 'ʐ'), ('кӏ', 'қ'),
     ('кхъ', 'qҳ'), ('къ', 'q'), ('лъ', 'ɬ'), ('лӏ', 'ԯ'), ('пӏ', 'ԥ'),
@@ -22,55 +20,55 @@ replacements = [
     ('щӏ', 'ɕ'), ('я', 'йа')
 ]
 
+# Reverse replacements for transcription
 reverse_replacements = {v: k for k, v in replacements}
 reverse_pattern = re.compile('|'.join(re.escape(key) for key in reverse_replacements))
 
 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 preprocess_audio(audio_tensor, original_sample_rate):
+    audio_tensor = audio_tensor.to(dtype=torch.float32)
+    audio_tensor = torch.mean(audio_tensor, dim=0, keepdim=True)  # Convert to mono
+    audio_tensor = audio_tensor / torch.max(torch.abs(audio_tensor))  # Normalize
+    audio_tensor = torchaudio.functional.resample(audio_tensor, orig_freq=original_sample_rate, new_freq=16000)  # Resample
+    return audio_tensor
 
-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))
-
-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)
+    return torch.tensor(filtered_audio, dtype=audio_tensor.dtype)
 
 @spaces.GPU
 def transcribe_speech(audio, progress=gr.Progress()):
-    if audio is None:  # Handle the NoneType error for microphone input
+    if audio is None:
         return "No audio received."
-
     progress(0.5, desc="Transcribing audio...")
-    transcription = pipe(audio, chunk_length_s=10)['text']
-
+    audio_np = audio.numpy().squeeze()
+    transcription = pipe(audio_np, chunk_length_s=10)['text']
     return replace_symbols_back(transcription)
 
 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)
-
-        progress(0.4, "Applying Wiener filter...")
-        apply_wiener_filter(audio_path)
-
-        progress(0.6, "Resampling audio...")
-        resample_audio(audio_path)
-
-    progress(0.8, "Transcribing audio...")
-    transcription = transcribe_speech(audio_path)
-
-    os.remove(audio_path)
+    
+    yt = YouTube(url)
+    stream = yt.streams.filter(only_audio=True).first()
+    audio_data = BytesIO()
+    stream.stream_to_buffer(audio_data)
+    audio_data.seek(0)
+    
+    try:
+        audio, original_sample_rate = torchaudio.load(audio_data)
+        audio = preprocess_audio(audio, original_sample_rate)
+
+        if apply_improvements:
+            progress(0.4, "Applying Wiener filter...")
+            audio = apply_wiener_filter(audio)
+
+        transcription = transcribe_speech(audio)
+
+    except Exception as e:
+        return str(e)
 
     return transcription