Update app.py

app.py

@@ -47,6 +47,11 @@ vision_model = AutoModelForCausalLM.from_pretrained(
 
 vision_processor = AutoProcessor.from_pretrained(VISION_MODEL_ID, trust_remote_code=True)
 
+# Helper functions
+# Initialize Parler-TTS
+tts_model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-mini-v1").to(device)
+tts_tokenizer = AutoTokenizer.from_pretrained("parler-tts/parler-tts-mini-v1")
+
 # Helper functions
 @spaces.GPU
 def stream_text_chat(message, history, system_prompt, temperature=0.8, max_new_tokens=1024, top_p=1.0, top_k=20):
@@ -77,9 +82,30 @@ def stream_text_chat(message, history, system_prompt, temperature=0.8, max_new_t
         thread.start()
 
     buffer = ""
+    audio_files = []
     for new_text in streamer:
         buffer += new_text
-        yield history + [[message, buffer]]
+        
+        # Generate speech for the new text
+        tts_input_ids = tts_tokenizer(new_text, return_tensors="pt").input_ids.to(device)
+        tts_description = "A clear and natural voice reads the text with moderate speed and expression."
+        tts_description_ids = tts_tokenizer(tts_description, return_tensors="pt").input_ids.to(device)
+        
+        with torch.no_grad():
+            audio_generation = tts_model.generate(input_ids=tts_description_ids, prompt_input_ids=tts_input_ids)
+        
+        audio_arr = audio_generation.cpu().numpy().squeeze()
+        
+        # Save the audio to a temporary file
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio:
+            sf.write(temp_audio.name, audio_arr, tts_model.config.sampling_rate)
+            audio_files.append(temp_audio.name)
+        
+        yield history + [[message, buffer]], audio_files
+
+    # Clean up temporary audio files
+    for audio_file in audio_files:
+        os.remove(audio_file)
 
 @spaces.GPU
 def process_vision_query(image, text_input):
@@ -106,6 +132,7 @@ def process_vision_query(image, text_input):
     response = vision_processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
     return response
 
+
 # Custom CSS
 custom_css = """
 body { background-color: #0b0f19; color: #e2e8f0; font-family: 'Arial', sans-serif;}
@@ -169,6 +196,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Base().set(
     with gr.Tab("Text Model (Phi-3.5-mini)"):
         chatbot = gr.Chatbot(height=400)
         msg = gr.Textbox(label="Message", placeholder="Type your message here...")
+        audio_output = gr.Audio(label="Generated Speech", autoplay=True)
         with gr.Accordion("Advanced Options", open=False):
             system_prompt = gr.Textbox(value="You are a helpful assistant", label="System Prompt")
             temperature = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.8, label="Temperature")
@@ -179,7 +207,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Base().set(
         submit_btn = gr.Button("Submit", variant="primary")
         clear_btn = gr.Button("Clear Chat", variant="secondary")
 
-        submit_btn.click(stream_text_chat, [msg, chatbot, system_prompt, temperature, max_new_tokens, top_p, top_k], [chatbot])
+        submit_btn.click(stream_text_chat, [msg, chatbot, system_prompt, temperature, max_new_tokens, top_p, top_k], [chatbot, audio_output])
         clear_btn.click(lambda: None, None, chatbot, queue=False)
 
     with gr.Tab("Vision Model (Phi-3.5-vision)"):