Added one more fix for image issue

Gradio_UI.py

@@ -61,6 +61,10 @@ def pull_messages_from_step(
             else:
                 content = str(args).strip()
 
+            # If the content looks like an image file path, convert it to Markdown so Gradio shows the image
+            if isinstance(content, str) and content.lower().endswith(('.png', '.jpg', '.jpeg', '.gif')):
+                content = f"![generated image]({content})"
+
             if used_code:
                 # Clean up the content by removing any end code tags
                 content = re.sub(r"```.*?\n", "", content)  # Remove existing code blocks
@@ -143,27 +147,23 @@ def stream_to_gradio(
             final_answer_step = step_log
             continue # Don't display the final answer step itself, process it at the end
 
-        # FIX: Check if token counts exist before trying to use them
-        if hasattr(agent.model, "last_input_token_count") and agent.model.last_input_token_count is not None:
-            if isinstance(step_log, ActionStep):
-                step_log.input_token_count = agent.model.last_input_token_count
-                step_log.output_token_count = agent.model.last_output_token_count
-
         for message in pull_messages_from_step(
             step_log,
         ):
             yield message
     
-    # Process the final answer correctly after the loop
+    # 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)
         
         # Directly check if the final answer is a path to an image file
         if isinstance(final_answer_value, str) and final_answer_value.lower().endswith(('.png', '.jpg', '.jpeg', '.gif')):
+            # Display the image inline using Markdown so Gradio will render it instead of plain text
+            markdown_img = f"![generated image]({final_answer_value})"
             yield gr.ChatMessage(
                 role="assistant",
-                content=final_answer_value, # Pass the file path directly to Gradio
+                content=markdown_img,
             )
         else:
             # Fallback for text or any other type of answer