Update app.py

app.py

@@ -8,6 +8,22 @@ import edge_tts
 import asyncio
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
+MAX_MAX_NEW_TOKENS = 2048
+DEFAULT_MAX_NEW_TOKENS = 1024
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+model_id = "prithivMLmods/FastThink-0.5B-Tiny"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+)
+model.eval()
+
+
 DESCRIPTION = """
 # QwQ Edge 💬
 """
@@ -26,25 +42,21 @@ h1 {
 }
 '''
 
-MAX_MAX_NEW_TOKENS = 2048
-DEFAULT_MAX_NEW_TOKENS = 1024
-MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-model_id = "prithivMLmods/FastThink-0.5B-Tiny"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    device_map="auto",
-    torch_dtype=torch.bfloat16,
-)
-model.eval()
+# List of voices
+voices = [
+    "en-US-JennyNeural",  # @tts1
+    "en-US-GuyNeural",    # @tts2
+    "en-US-AriaNeural",   # @tts3
+    "en-US-DavisNeural",  # @tts4
+    "en-US-JaneNeural",   # @tts5
+    "en-US-JasonNeural",  # @tts6
+    "en-US-NancyNeural",  # @tts7
+    "en-US-TonyNeural",   # @tts8
+]
 
 
-async def text_to_speech(text: str, output_file="output.mp3"):
+async def text_to_speech(text: str, voice: str, output_file="output.mp3"):
     """Convert text to speech using Edge TTS and save as MP3"""
-    voice = "en-US-GuyNeural"  # Change this to your preferred voice
     communicate = edge_tts.Communicate(text, voice)
     await communicate.save(output_file)
     return output_file
@@ -62,7 +74,24 @@ def generate(
 ):
     """Generates chatbot response and handles TTS requests"""
     is_tts = message.strip().lower().startswith("@tts")
-    message = message.replace("@tts", "").strip()
+    tts_index = None
+
+    if is_tts:
+        # Extract the number after @tts
+        tts_part = message.strip().lower().split()[0]  # Get the @ttsX part
+        if len(tts_part) > 4:  # Check if it's @ttsX (e.g., @tts1, @tts2, etc.)
+            try:
+                tts_index = int(tts_part[4:]) - 1  # Convert to 0-based index
+                if tts_index < 0 or tts_index >= len(voices):
+                    gr.Warning(f"Invalid TTS voice index. Using default voice.")
+                    tts_index = 0
+            except ValueError:
+                gr.Warning(f"Invalid TTS voice index. Using default voice.")
+                tts_index = 0
+        else:
+            tts_index = 0  # Default to the first voice if no number is provided
+
+        message = message.replace(tts_part, "").strip()  # Remove @ttsX from the message
 
     conversation = [*chat_history, {"role": "user", "content": message}]
 
@@ -95,7 +124,8 @@ def generate(
     final_response = "".join(outputs)
 
     if is_tts:
-        output_file = asyncio.run(text_to_speech(final_response))
+        voice = voices[tts_index]  # Select the voice based on the index
+        output_file = asyncio.run(text_to_speech(final_response, voice))
         yield gr.Audio(output_file, autoplay=True)  # Return playable audio
     else:
         yield final_response  # Return text response
@@ -112,12 +142,12 @@ demo = gr.ChatInterface(
     ],
     stop_btn=None,
     examples=[
-        ["@tts Who is Nikola Tesla, and why did he die?"],
-        ["A train travels 60 kilometers per hour. If it travels for 5 hours, how far will it travel in total?"],
+        ["@tts1 Who is Nikola Tesla, and why did he die?"],
+        ["@tts2 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."],
-        ["@tts What causes rainbows to form?"],
+        ["@tts3 What causes rainbows to form?"],
         ["Rewrite the following sentence in passive voice: 'The dog chased the cat.'"],
-        ["@tts What is the capital of France?"],
+        ["@tts4 What is the capital of France?"],
     ],
     cache_examples=False,
     type="messages",