Update

app.py

@@ -173,7 +173,9 @@ def visualize_attention(
 
         # Convert plot to image
         fig.canvas.draw()
-        vis_image = Image.frombytes('RGB', fig.canvas.get_width_height(), fig.canvas.buffer_rgba())
+        data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+        data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+        vis_image = Image.fromarray(data)
         visualizations.append(vis_image)
         plt.close(fig)
 
@@ -200,11 +202,14 @@ def visualize_attention(
 
     # Convert plot to image
     fig.canvas.draw()
-    rollout_image = Image.frombytes('RGB', fig.canvas.get_width_height(), fig.canvas.buffer_rgba())
+    data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    rollout_image = Image.fromarray(data)
     plt.close(fig)
 
     return visualizations, rollout_image
 
+
 # Create Gradio interface
 iface = gr.Interface(
     fn=visualize_attention,