Update app.py

app.py

@@ -6,6 +6,8 @@ import matplotlib.pyplot as plt
 import seaborn as sns
 import numpy as np
 import gradio as gr
+from io import BytesIO
+from PIL import Image
 
 def get_heatmap(sequence):
   logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)
@@ -64,11 +66,15 @@ def get_heatmap(sequence):
   plt.yticks(rotation=0)
   plt.xticks(y_tick_positions, y_tick_labels, rotation = 0)
 
-  fig = plt.gcf()
-  plt.close(fig)
-
-  return fig
-
+# Save the figure to a BytesIO object
+  buf = BytesIO()
+  plt.savefig(buf, format='png')
+  buf.seek(0)
+  plt.close()
+    
+# Convert BytesIO object to an image
+  img = Image.open(buf)
+  return img
 
 
 demo = gr.Interface(fn=get_heatmap, inputs="text", outputs="image")