Fix the image object

Gradio_UI.py

@@ -159,31 +159,39 @@ def stream_to_gradio(
             yield message
     
     # FIX: Process the final answer correctly after the loop
-    if final_answer_step:
-        # Extract the actual value from the FinalAnswerStep object
-        final_answer_value = getattr(final_answer_step, 'final_answer', final_answer_step)
-        
-        # Convert the value (e.g., a file path string) into the correct Agent type (e.g., AgentImage)
+if final_answer_step:
+    # Extract the actual value from the FinalAnswerStep object
+    final_answer_value = getattr(final_answer_step, 'final_answer', final_answer_step)
+    
+    # Check if the final answer is a string that looks like an image file path
+    if isinstance(final_answer_value, str) and final_answer_value.endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp')):
+        # It's an image file path, create an AgentImage object
+        from smolagents.agent_types import AgentImage
+        processed_answer = AgentImage(value=final_answer_value)
+    else:
+        # Try the default conversion
         processed_answer = handle_agent_output_types(final_answer_value)
 
-        if isinstance(processed_answer, AgentText):
-            yield gr.ChatMessage(
-                role="assistant",
-                content=f"**Final answer:**\n{processed_answer.to_string()}\n",
-            )
-        elif isinstance(processed_answer, AgentImage):
-            yield gr.ChatMessage(
-                role="assistant",
-                content=processed_answer.to_string(), # Gradio automatically handles file paths
-            )
-        elif isinstance(processed_answer, AgentAudio):
-            yield gr.ChatMessage(
-                role="assistant",
-                content=processed_answer.to_string(), # Gradio automatically handles file paths
-            )
-        else:
-            # Fallback for any other type
-            yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(processed_answer)}")
+    if isinstance(processed_answer, AgentText):
+        yield gr.ChatMessage(
+            role="assistant",
+            content=f"**Final answer:**\n{processed_answer.to_string()}\n",
+        )
+    elif isinstance(processed_answer, AgentImage):
+        # For images, Gradio needs the file path directly
+        image_path = processed_answer.to_string()
+        yield gr.ChatMessage(
+            role="assistant",
+            content={"file": image_path, "alt_text": "Generated image"},
+        )
+    elif isinstance(processed_answer, AgentAudio):
+        yield gr.ChatMessage(
+            role="assistant",
+            content={"file": processed_answer.to_string(), "alt_text": "Generated audio"},
+        )
+    else:
+        # Fallback for any other type
+        yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(processed_answer)}")
 
 
 class GradioUI: