Update app.py

app.py

@@ -3,37 +3,38 @@ import gradio as gr
 import os
 from project_model import process_inputs, session
 
-# --- Handle Initial Upload ---
-def handle_initial(image, audio):
-    if image is None or audio is None:
-        return "❗ Please upload both an image and an audio clip.", None
-
-    message, answer_audio = process_inputs(session, image=image, audio_path=audio)
-
-    # Save uploaded image
-    image_save_path = "uploaded_image.png"
-    image.save(image_save_path)
-
-    # Markdown text
-    markdown_reply = f"**{message}**\n\n![Context Image](file/{image_save_path})"
-
-    return markdown_reply, answer_audio
-
-# --- Handle Follow-up ---
-def handle_followup(followup_audio):
-    if followup_audio is None:
-        return "❗ Please record a follow-up question.", None
-
-    message, answer_audio = process_inputs(session, audio_path=followup_audio)
-
-    # Reuse saved image
-    image_save_path = "uploaded_image.png"
-    session.current_image.save(image_save_path)
-
-    # Markdown text
-    markdown_reply = f"**{message}**\n\n![Context Image](file/{image_save_path})"
-
-    return markdown_reply, answer_audio
+# --- Handle both initial upload and follow-up ---
+def handle_submit(image, audio, followup_audio):
+    """
+    Handles both initial question and follow-up.
+    Priority:
+    - If new image + audio => initial upload
+    - If only followup audio => follow-up
+    """
+    if image is not None and audio is not None:
+        # Initial case
+        message, answer_audio = process_inputs(session, image=image, audio_path=audio)
+
+        # Save uploaded image
+        image_save_path = "uploaded_image.png"
+        image.save(image_save_path)
+
+        markdown_reply = f"**{message}**\n\n![Context Image](file/{image_save_path})"
+        return markdown_reply, answer_audio
+
+    elif followup_audio is not None:
+        # Follow-up case
+        message, answer_audio = process_inputs(session, audio_path=followup_audio)
+
+        # Reuse saved image
+        image_save_path = "uploaded_image.png"
+        session.current_image.save(image_save_path)
+
+        markdown_reply = f"**{message}**\n\n![Context Image](file/{image_save_path})"
+        return markdown_reply, answer_audio
+
+    else:
+        return "❗ Please upload image/audio or record a follow-up.", None
 
 # --- Gradio App ---
 with gr.Blocks() as demo:
@@ -43,25 +44,19 @@ with gr.Blocks() as demo:
         with gr.Column():
             image_input = gr.Image(label="Upload or Capture Image", sources=["upload", "webcam"], type="pil")
             audio_input = gr.Audio(label="Initial Question (Voice)", sources=["microphone"], type="filepath")
-            submit_btn = gr.Button("Submit Initial Q&A")
+            gr.Markdown("### 🎤 OR Ask a Follow-up Question")
+            followup_audio_input = gr.Audio(label="Follow-up Question (Voice)", sources=["microphone"], type="filepath")
 
-            gr.Markdown("### 🎤 Ask a Follow-up Question")
-            followup_audio_input = gr.Audio(label="Follow-up Question", sources=["microphone"], type="filepath")
-            followup_btn = gr.Button("Ask Follow-up")
+            submit_btn = gr.Button("Submit Question / Follow-up")  # <== 🔥 Single button
 
         with gr.Column():
-            status_output = gr.Markdown(label="Response")  # <== 🔥 Changed to Markdown
+            status_output = gr.Markdown(label="Response")
             audio_output = gr.Audio(label="🔊 Response Audio", interactive=False)
 
+    # Hook up the single submit button
     submit_btn.click(
-        fn=handle_initial,
-        inputs=[image_input, audio_input],
-        outputs=[status_output, audio_output]
-    )
-
-    followup_btn.click(
-        fn=handle_followup,
-        inputs=[followup_audio_input],
+        fn=handle_submit,
+        inputs=[image_input, audio_input, followup_audio_input],
         outputs=[status_output, audio_output]
     )