Update app.py

app.py

@@ -1,23 +1,35 @@
+import torch
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, VitsModel, AutoTokenizer
 
 asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-tiny")
+
 llm_pipe = pipeline(
     "text-generation",
     model="microsoft/Phi-3-mini-4k-instruct",
-    device_map="auto",      # needs `pip install accelerate`; remove to run on CPU only
-    torch_dtype="auto"
+    device_map="auto",
+    torch_dtype="auto",
 )
-tts_pipe = pipeline("text-to-speech", model="facebook/mms-tts-eng")
+
+tts_model = VitsModel.from_pretrained("facebook/mms-tts-eng")
+tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
+
+
+def text_to_speech(text):
+    inputs = tts_tokenizer(text, return_tensors="pt")
+    with torch.no_grad():
+        output = tts_model(**inputs).waveform
+    audio = output.squeeze().cpu().numpy()
+    sampling_rate = tts_model.config.sampling_rate
+    return sampling_rate, audio
+
 
 def process_audio(audio_path):
     if audio_path is None:
         return "Please record something first.", None
 
-    # 1. Speech -> text
     user_text = asr_pipe(audio_path)["text"].strip()
 
-    # 2. Text -> reply (chat format, reply only, capped length)
     messages = [
         {"role": "system", "content": "You are a helpful voice assistant. Answer in 1-3 short sentences."},
         {"role": "user", "content": user_text},
@@ -27,14 +39,11 @@ def process_audio(audio_path):
         max_new_tokens=100,
         do_sample=False,
     )
-    # Chat-formatted input returns the full conversation in "generated_text";
-    # the assistant's reply is the last message in that list.
     reply = output[0]["generated_text"][-1]["content"].strip()
 
-    # 3. Reply -> speech
-    speech = tts_pipe(reply)
-    audio = speech["audio"].squeeze()
-    return reply, (speech["sampling_rate"], audio)
+    sampling_rate, audio = text_to_speech(reply)
+    return reply, (sampling_rate, audio)
+
 
 demo = gr.Interface(
     fn=process_audio,