Update Gradio_UI.py

Gradio_UI.py

@@ -162,9 +162,24 @@ def stream_to_gradio(
             content=f"**Final answer:**\n{final_answer.to_string()}\n",
         )
     elif isinstance(final_answer, AgentImage):
+        img_path = final_answer.to_string()
+    
+        # If the returned string isn't a valid file path, save the image ourselves
+        if not isinstance(img_path, str) or not os.path.exists(img_path):
+            os.makedirs("generated_images", exist_ok=True)
+            img_path = os.path.join("generated_images", "generated_image.png")
+    
+            # Try to save AgentImage.value (usually PIL image)
+            try:
+                final_answer.value.save(img_path)
+            except Exception:
+                # Fallback for raw bytes
+                with open(img_path, "wb") as f:
+                    f.write(final_answer.value)
+    
         yield gr.ChatMessage(
             role="assistant",
-            content={"path": final_answer.to_string(), "mime_type": "image/png"},
+            content={"path": img_path, "mime_type": "image/png"},
         )
     elif isinstance(final_answer, AgentAudio):
         yield gr.ChatMessage(