Update app.py

app.py

@@ -88,4 +88,23 @@ agent = CodeAgent(
 )
 
 
-GradioUI(agent).launch()
+# -------------------------
+# Launch Gradio UI
+# -------------------------
+# Custom wrapper to automatically unwrap BytesIO images from FinalAnswerStep
+def gradio_output_wrapper(agent_output):
+    """
+    Checks if agent_output.final_answer is a BytesIO image and returns it directly
+    for Gradio to display.
+    """
+    try:
+        # If agent_output has final_answer attribute (FinalAnswerStep)
+        buffer = getattr(agent_output, "final_answer", None)
+        if isinstance(buffer, BytesIO):
+            buffer.seek(0)
+            return buffer  # Gradio can render BytesIO as an image
+        return agent_output  # fallback for text or other outputs
+    except Exception:
+        return agent_output
+
+GradioUI(agent, output_wrapper=gradio_output_wrapper).launch()