Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, VitsForConditionalGeneration, VitsProcessor
+from transformers import AutoTokenizer, AutoModelForCausalLM, VitsModel
 from nemo.collections.asr.models import ASRModel
 
 
@@ -27,9 +27,9 @@ proc_model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True, 
 )
 
-# Load the TTS model and processor
-tts_processor = VitsProcessor.from_pretrained("facebook/mms-tts-eng")
-tts_model = VitsForConditionalGeneration.from_pretrained("facebook/mms-tts-eng")
+# Load the TTS model
+tts_model = VitsModel.from_pretrained("facebook/mms-tts-eng")
+tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
 
 
 def process_speech(speech):
@@ -39,14 +39,12 @@ def process_speech(speech):
     # Process the text
     inputs = proc_tokenizer.encode(transcription + proc_tokenizer.eos_token, return_tensors='pt')
     outputs = proc_model.generate(inputs, max_length=100, temperature=0.7, pad_token_id=proc_tokenizer.eos_token_id)
-    processed_text = proc_tokenizer.decode(outputs[0], skip_special_tokens=True)
-
+    text = proc_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    processed_text = tts_tokenizer(text, return_tensors="pt")
+    
     # Convert the processed text to speech
-    inputs = tts_processor(processed_text, return_tensors="pt")
     with torch.no_grad():
-        logits = tts_model(inputs["input_ids"]).logits
-    predicted_ids = torch.argmax(logits, dim=-1)
-    audio = tts_processor.decode(predicted_ids)
+        audio = tts_model(**inputs).waveform   
 
     return audio