Update app.py

app.py

@@ -88,23 +88,4 @@ agent = CodeAgent(
 )
 
 
-# -------------------------
-# 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()
+GradioUI(agent).launch()