Update tts.py

tts.py

@@ -1,21 +1,18 @@
 import torch
-from transformers import SpeechT5ForTextToSpeech, SpeechT5Processor
-import sentencepiece
+from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech
+from datasets import load_dataset
+import soundfile as sf
 
 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):
-    if not text:
-        return "ERROR: Please provide text for synthesis"
-    
-    inputs = processor(text, return_tensors="pt")
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-    inputs = inputs.to(device)
+def synthesize_speech(text_input):
+    inputs = processor(text=text_input, return_tensors="pt")
 
     with torch.no_grad():
-        speech = model.generate(**inputs)
+        speech = model.generate_speech(inputs["input_ids"], vocoder=vocoder)
 
-    return processor.decode(speech, skip_special_tokens=True)
+    sf.write("output.wav", speech.numpy(), 16000)
+    return "output.wav"