Update app.py

app.py

@@ -1,41 +1,25 @@
 from svoice.separate import *
-import scipy.io as sio
-from scipy.io.wavfile import write
+import scipy.io.wavfile as wav
 import gradio as gr
 import os
-# from transformers import AutoProcessor, pipeline
-# from optimum.onnxruntime import ORTModelForSpeechSeq2Seq
-import whisper
+import torch
+import soundfile as sf
+from transformers import pipeline
 from glob import glob
 load_model()
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
 BASE_PATH = os.path.dirname(os.path.abspath(__file__))
 os.makedirs('input', exist_ok=True)
 os.makedirs('separated', exist_ok=True)
 
-# print("Loading ASR model...")
-# processor = AutoProcessor.from_pretrained("openai/whisper-small")
-# if not os.path.exists("whisper_checkpoint"):
-#     model = ORTModelForSpeechSeq2Seq.from_pretrained("openai/whisper-small", from_transformers=True)
-#     speech_recognition_pipeline = pipeline(
-#     "automatic-speech-recognition",
-#         model=model,
-#         feature_extractor=processor.feature_extractor,
-#         tokenizer=processor.tokenizer,
-#     )
-#     os.makedirs('whisper_checkpoint', exist_ok=True)
-#     model.save_pretrained("whisper_checkpoint")
-# else:
-#     model = ORTModelForSpeechSeq2Seq.from_pretrained("whisper_checkpoint", from_transformers=False)
-#     speech_recognition_pipeline = pipeline(
-#     "automatic-speech-recognition",
-#         model=model,
-#         feature_extractor=processor.feature_extractor,
-#         tokenizer=processor.tokenizer,
-#     )
-# print("Whisper ASR model loaded.")
+print("Loading ASR model...")
+pipe = pipeline("automatic-speech-recognition", model="openai/whisper-base", device=0 if device == "cuda" else -1)
+print("ASR model loaded!")
 
-model = whisper.load_model("base")
+def transcribe_audio(audiopath):
+    audio_input, sr = sf.read(audiopath)
+    return pipe(audio_input, sampling_rate=sr, return_tensors=False, padding=True, max_new_tokens=500)['text']
 
 def separator(audio, rec_audio, example):
     outputs= {}
@@ -44,21 +28,22 @@ def separator(audio, rec_audio, example):
     for f in glob('separated/*'):
         os.remove(f)
     if audio:
-        write('input/original.wav', audio[0], audio[1])
+        wav.write('input/original.wav', audio[0], audio[1])
     elif rec_audio:
-        write('input/original.wav', rec_audio[0], rec_audio[1])
+        wav.write('input/original.wav', rec_audio[0], rec_audio[1])
     else:
         os.system(f'cp {example} input/original.wav')
     separate_demo(mix_dir="./input")
     separated_files = glob(os.path.join('separated', "*.wav"))
-    separated_files = [f for f in separated_files if "original.wav" not in f]
-    outputs['transcripts'] = []
-    for file in sorted(separated_files):
-        # separated_audio = sio.wavfile.read(file)
-        # outputs['transcripts'].append(speech_recognition_pipeline(separated_audio[1])['text'])
-        outputs['transcripts'].append(whisper.transcribe(file)["text"])
-    return sorted(separated_files) + outputs['transcripts']
+    separated_files = sorted([f for f in separated_files if "original.wav" not in f])
+    outputs["transcripts"] = []
     
+    for i, f in enumerate(separated_files):
+        print(f"Transcribing separated audio {i+1} ...")
+        outputs["transcripts"].append(transcribe_audio(f))
+        print("Text:", outputs["transcripts"][-1])
+    return separated_files + outputs['transcripts']
+
 def set_example_audio(example: list) -> dict:
     return gr.Audio.update(value=example[0])
 
@@ -108,7 +93,7 @@ with demo:
         "samples/mixture2.wav",
         "samples/mixture3.wav"
     ]
-    example_selector = gr.inputs.Dropdown(examples, label="Example Audio", default="samples/mixture1.wav")
+    example_selector = gr.inputs.Radio(examples, label="Example Audio")
     button.click(separator, inputs=[input_audio, rec_audio, example_selector], outputs=outputs_audio + outputs_text)
 
 demo.launch()