Update app.py

app.py

@@ -1,16 +1,17 @@
 import os
 from collections.abc import Iterator
 from threading import Thread
+import ai_gradio
 import gradio as gr
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 DESCRIPTION = """
-# QwQ Tiny
+# QwQ Tiny & Kokoro Chatbot
 """
 
-css ='''
+css = '''
 h1 {
   text-align: center;
   display: block;
@@ -39,6 +40,12 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
+# Load Kokoro chatbot
+kokoro_demo = gr.load(
+    name='onnx-community:kokoro-v1_0',
+    src='ai_gradio.registry',
+)
+kokoro_demo.fn = spaces.GPU()(kokoro_demo.fn)
 
 @spaces.GPU
 def generate(
@@ -50,6 +57,21 @@ def generate(
     top_k: int = 50,
     repetition_penalty: float = 1.2,
 ) -> Iterator[str]:
+    # Check for tags in the message
+    if "@voice" in message:
+        # Remove the tag from the message
+        message = message.replace("@voice", "").strip()
+        # Use Kokoro for voice
+        yield kokoro_demo.fn(message)
+    elif "@text" in message:
+        # Remove the tag from the message
+        message = message.replace("@text", "").strip()
+        # Fall through to text generation
+    else:
+        # Default to text generation if no tag is specified
+        pass
+
+    # Text generation logic (QwQ Tiny)
     conversation = [*chat_history, {"role": "user", "content": message}]
 
     input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
@@ -78,7 +100,6 @@ def generate(
         outputs.append(text)
         yield "".join(outputs)
 
-
 demo = gr.ChatInterface(
     fn=generate,
     additional_inputs=[
@@ -124,6 +145,8 @@ demo = gr.ChatInterface(
         ["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.'"],
+        ["@voice Hello, how are you today?"],
+        ["@text Explain the theory of relativity."],
     ],
     cache_examples=False,
     type="messages",
@@ -132,6 +155,5 @@ demo = gr.ChatInterface(
     fill_height=True,
 )
 
-
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=20).launch()