Updated the UI

app.py

@@ -1,63 +1,56 @@
 import gradio as gr
+import time
 from src.pipeline import generate_report
 
 # ------------------------------------------------------------------
-# 1. Pre-load models on Space start-up
+# 1. Pre-load models (unchanged)
 # ------------------------------------------------------------------
-print("Pre-loading models for fast inference …")
-try:
-    from src.tools_loader import get_tools          # downloads BiomedCLIP + SPECTER-2
-    _ = get_tools()
-    print("Models pre-loaded successfully!")
-except Exception as e:
-    print(f"Model pre-loading failed: {e}")
+from src.tools_loader import get_tools
+_ = get_tools()
 
 # ------------------------------------------------------------------
-# 2. Inference wrapper
+# 2. Helper: streaming generator
 # ------------------------------------------------------------------
 def process_upload(image_path: str):
-    """Run the multi-agent pipeline on an uploaded chest X-ray."""
+    """
+    Streamed generator: yields placeholder first,
+    then final report with inference time.
+    Gradio automatically shows a spinner / progress bar
+    while the function is running.
+    """
     if image_path is None:
-        return "Please upload a chest X-ray image."
+        yield "Please upload a chest-X-ray image."
+        return
 
-    try:
-        report = generate_report(image_path)
-        return report
-    except Exception as e:
-        return f"Error processing image: {e}"
+    start = time.time()
+    # Initial placeholder so the textarea updates immediately
+    yield " **Generating report...**\n\nThis may take a few seconds..."
+
+    # Optional manual progress bar
+    # with gr.Progress(track_tqdm=True) as progress:
+    #     report = generate_report(image_path, progress=progress)
+
+    report = generate_report(image_path)
+
+    elapsed = time.time() - start
+    yield f"### Radiology Report\n{report}\n\n---\n*Generated in `{elapsed:0.1f}` s*"
 
 # ------------------------------------------------------------------
 # 3. Gradio UI
 # ------------------------------------------------------------------
-with gr.Blocks(title="Multi-Agent Radiology Assistant") as demo:
-    gr.Markdown(
-        """
-        # Multi-Agent Radiology Assistant  
-        Upload a chest X-ray and receive an AI-generated report produced by a multi-agent pipeline.
-        """
-    )
+with gr.Blocks() as demo:
+    gr.Markdown("# Multi-Agent Radiology Assistant")
+    with gr.Row():
+        input_image = gr.Image(type="filepath", label="Upload Chest X-ray", height=400)
+        output_report = gr.Markdown()
 
-    # --- Upload widget + button ------------------------------------------------
-    with gr.Column():
-        input_image = gr.Image(
-            type="filepath",
-            label="Upload Chest X-ray",
-            height=400
-        )
-        process_btn = gr.Button("Generate Report", variant="primary")
-
-    # --- Report output ---------------------------------------------------------
-    output_report = gr.Markdown(label="Radiology Report", show_label=True)
-
-    # --- Wire everything together ---------------------------------------------
-    process_btn.click(
-        fn=process_upload,
+    generate_btn = gr.Button("Generate Report")
+
+    generate_btn.click(
+        fn=process_upload,            # generator function
         inputs=input_image,
         outputs=output_report
     )
 
-    gr.Markdown("### Need an example?  \nUse any frontal CXR PNG file and click **Generate Report**.")
-
-# ------------------------------------------------------------------
 if __name__ == "__main__":
     demo.launch()