Update tts.py

tts.py

@@ -1,18 +1,17 @@
 import torch
-from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech
-from datasets import load_dataset
-import soundfile as sf
+from transformers import SpeechT5ForTextToSpeech, SpeechT5Processor
 
 MODEL_ID = "microsoft/speecht5_tts"
 processor = SpeechT5Processor.from_pretrained(MODEL_ID)
 model = SpeechT5ForTextToSpeech.from_pretrained(MODEL_ID)
-vocoder = torch.hub.load("snakers4/silero-vad", "silero_vad", force_reload=True)
 
-def synthesize_speech(text_input):
-    inputs = processor(text=text_input, return_tensors="pt")
+def synthesize_speech(text):
+    inputs = processor(text, return_tensors="pt")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    inputs = inputs.to(device)
 
     with torch.no_grad():
-        speech = model.generate_speech(inputs["input_ids"], vocoder=vocoder)
+        speech = model.generate(**inputs)
 
-    sf.write("output.wav", speech.numpy(), 16000)
-    return "output.wav"
+    return processor.decode(speech, skip_special_tokens=True)