Use custom Blocks interface with dedicated image output for grounding visualizations

app.py

@@ -66,6 +66,7 @@ agent = Agent(
 print(f"Tools loaded: {len(tools)}")
 
 import glob
+from PIL import Image
 
 def chat(message, history):
     config = {"configurable": {"thread_id": "default"}}
@@ -88,23 +89,60 @@ def chat(message, history):
     assistant_message = response["messages"][-1].content
 
     # Check for grounding visualization images
+    viz_image = None
     viz_files = glob.glob("temp/grounding_*.png")
     if viz_files:
         # Get the most recent visualization
         viz_files.sort(key=os.path.getmtime, reverse=True)
         latest_viz = viz_files[0]
+        try:
+            viz_image = Image.open(latest_viz)
+        except:
+            pass
 
-        # Append image reference that Gradio can render
-        assistant_message += f"\n\n**Grounding Visualization:**\n\n<img src='/file={latest_viz}' style='max-width: 600px;'/>"
+    return assistant_message, viz_image
 
-    return assistant_message
+# Custom interface with image output
+with gr.Blocks() as demo:
+    gr.Markdown(f"# MedRAX2 - Medical AI Assistant\n**Device:** {device} | **Tools:** {len(tools)} loaded")
 
-demo = gr.ChatInterface(
-    fn=chat,
-    title="MedRAX2 - Medical AI Assistant",
-    description=f"Device: {device} | Tools: {len(tools)} loaded",
-    multimodal=True
-)
+    chatbot = gr.Chatbot(type="messages")
+    viz_output = gr.Image(label="Grounding Visualization", visible=True)
+
+    msg = gr.MultimodalTextbox(
+        label="Message",
+        placeholder="Upload an X-ray image and ask a question...",
+        file_types=["image"]
+    )
+
+    def respond(message, chat_history):
+        # Get response and visualization
+        bot_message, viz_image = chat(message, chat_history)
+
+        # Add to chat history
+        if isinstance(message, dict):
+            text = message.get("text", "")
+            files = message.get("files", [])
+            if files:
+                chat_history.append({"role": "user", "content": [{"type": "image", "image": files[0]}, {"type": "text", "text": text}]})
+            else:
+                chat_history.append({"role": "user", "content": text})
+        else:
+            chat_history.append({"role": "user", "content": message})
+
+        chat_history.append({"role": "assistant", "content": bot_message})
+
+        return "", chat_history, viz_image
+
+    msg.submit(respond, [msg, chatbot], [msg, chatbot, viz_output])
+
+    gr.Examples(
+        examples=[
+            [{"text": "What do you see in this X-ray?", "files": []}],
+            [{"text": "Can you show me where exactly using grounding?", "files": []}],
+        ],
+        inputs=msg,
+    )
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)