retry button

app.py

@@ -265,6 +265,9 @@ def _chat_tab(examples):
             interactive=True
         )
     
+    # Add visualization display
+    visualization_display = gr.Image(label="Audio Visualization", visible=True)
+    
     chatbot = gr.Chatbot(
         label="Model inputs",
         elem_id="chatbot",
@@ -279,16 +282,30 @@ def _chat_tab(examples):
     clear_button = gr.ClearButton(components=[chatbot, chat_input], visible=False)
 
     # Function to update audio visualization in chat
-    def update_chat_visualization(audio_path, view_type):
-        if not audio_path:
+    def update_chat_visualization(message, view_type):
+        if not message or not message.get("files"):
+            return None
+        
+        # Get the first audio file from the message
+        audio_files = message.get("files", [])
+        if not audio_files:
             return None
+            
+        audio_path = audio_files[0]
         return audio_visualization_component(audio_path, view_type.lower())
     
     # Connect visualization update for chat
     view_selector.change(
         update_chat_visualization,
         [chat_input, view_selector],
-        [chat_input]
+        [visualization_display]
+    )
+    
+    # Also update when audio is uploaded
+    chat_input.change(
+        update_chat_visualization,
+        [chat_input, view_selector],
+        [visualization_display]
     )
 
     chat_input.submit(add_message, [chatbot, chat_input], [chatbot, chat_input])