Update app.py

app.py

@@ -4,6 +4,8 @@ from threading import Thread
 import gradio as gr
 import spaces
 import torch
+import edge_tts
+import asyncio
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 DESCRIPTION = """
@@ -40,6 +42,14 @@ model = AutoModelForCausalLM.from_pretrained(
 model.eval()
 
 
+async def text_to_speech(text: str, output_file="output.mp3"):
+    """Convert text to speech using Edge TTS and save as MP3"""
+    voice = "en-US-JennyNeural"  # Change this to your preferred voice
+    communicate = edge_tts.Communicate(text, voice)
+    await communicate.save(output_file)
+    return output_file
+
+
 @spaces.GPU
 def generate(
     message: str,
@@ -49,7 +59,11 @@ def generate(
     top_p: float = 0.9,
     top_k: int = 50,
     repetition_penalty: float = 1.2,
-) -> Iterator[str]:
+):
+    """Generates chatbot response and handles TTS requests"""
+    is_tts = message.strip().lower().startswith("@tts")
+    message = message.replace("@tts", "").strip()
+
     conversation = [*chat_history, {"role": "user", "content": message}]
 
     input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
@@ -78,52 +92,31 @@ def generate(
         outputs.append(text)
         yield "".join(outputs)
 
+    final_response = "".join(outputs)
+
+    if is_tts:
+        output_file = asyncio.run(text_to_speech(final_response))
+        return output_file  # Return MP3 file
+
+    return final_response  # Return text response
+
 
 demo = gr.ChatInterface(
     fn=generate,
     additional_inputs=[
-        gr.Slider(
-            label="Max new tokens",
-            minimum=1,
-            maximum=MAX_MAX_NEW_TOKENS,
-            step=1,
-            value=DEFAULT_MAX_NEW_TOKENS,
-        ),
-        gr.Slider(
-            label="Temperature",
-            minimum=0.1,
-            maximum=4.0,
-            step=0.1,
-            value=0.6,
-        ),
-        gr.Slider(
-            label="Top-p (nucleus sampling)",
-            minimum=0.05,
-            maximum=1.0,
-            step=0.05,
-            value=0.9,
-        ),
-        gr.Slider(
-            label="Top-k",
-            minimum=1,
-            maximum=1000,
-            step=1,
-            value=50,
-        ),
-        gr.Slider(
-            label="Repetition penalty",
-            minimum=1.0,
-            maximum=2.0,
-            step=0.05,
-            value=1.2,
-        ),
+        gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS),
+        gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6),
+        gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9),
+        gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50),
+        gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2),
     ],
     stop_btn=None,
     examples=[
         ["A train travels 60 kilometers per hour. If it travels for 5 hours, how far will it travel in total?"],
-        ["Write a Python function to check if a number is prime. "],
+        ["Write a Python function to check if a number is prime."],
         ["What causes rainbows to form?"],
         ["Rewrite the following sentence in passive voice: 'The dog chased the cat.'"],
+        ["@tts What is the capital of France?"],
     ],
     cache_examples=False,
     type="messages",
@@ -132,6 +125,5 @@ demo = gr.ChatInterface(
     fill_height=True,
 )
 
-
 if __name__ == "__main__":
     demo.queue(max_size=20).launch()