Update Gradio_UI.py

Gradio_UI.py

@@ -130,8 +130,8 @@ def stream_to_gradio(
             total_input_tokens += agent.model.last_input_token_count
             total_output_tokens += agent.model.last_output_token_count
             if isinstance(step_log, ActionStep):
-                step_log.input_token_count = agent.model.last_input_token_count# or 0 # or 0 added for non-Gemini models
-                step_log.output_token_count = agent.model.last_output_token_count# or 0 # or 0 added for non-Gemini models
+                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,