Update app.py

app.py

@@ -50,11 +50,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
 # 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):
@@ -67,10 +67,12 @@ def stream_text_chat(message, history, system_prompt, temperature=0.8, max_new_t
     conversation.append({"role": "user", "content": message})
 
     input_ids = text_tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(text_model.device)
+    attention_mask = torch.ones_like(input_ids)  # Create attention mask
     streamer = TextIteratorStreamer(text_tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
 
     generate_kwargs = dict(
         input_ids=input_ids,
+        attention_mask=attention_mask,  # Pass attention mask
         max_new_tokens=max_new_tokens,
         do_sample=temperature > 0,
         top_p=top_p,
@@ -85,7 +87,7 @@ def stream_text_chat(message, history, system_prompt, temperature=0.8, max_new_t
         thread.start()
 
     buffer = ""
-    audio_files = []
+    audio_buffer = np.array([])
     for new_text in streamer:
         buffer += new_text
         
@@ -97,18 +99,10 @@ def stream_text_chat(message, history, system_prompt, temperature=0.8, max_new_t
         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()
+        new_audio = audio_generation.cpu().numpy().squeeze()
+        audio_buffer = np.concatenate((audio_buffer, new_audio))
         
-        # 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)
+        yield history + [[message, buffer]], (tts_model.config.sampling_rate, audio_buffer)
         
 @spaces.GPU
 def process_vision_query(image, text_input):
@@ -212,7 +206,6 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Base().set(
 
         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)"):
         with gr.Row():
             with gr.Column(scale=1):