Update app.py

app.py

@@ -2,14 +2,15 @@ import gradio as gr
 import numpy as np
 import torch
 from datasets import load_dataset
-from transformers import VitsModel, VitsTokenizer
+
 from transformers import pipeline
+from transformers import VitsModel, VitsTokenizer
 
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 # load speech translation checkpoint
-asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-tiny", device=device)
+asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-base", device=device)
 
 
 
@@ -27,8 +28,8 @@ def translate(audio):
 def synthesise(text):
     inputs = processor(text=text, return_tensors="pt")
     with torch.no_grad():
-      output = model(**inputs)
-    return output['audio']
+      speech = model(inputs["input_ids"].to(device))
+    return speech.audio[0]
 
 
 def speech_to_speech_translation(audio):