Update app.py

app.py

@@ -8,6 +8,8 @@ from matplotlib.patches import Patch
 from collections import defaultdict
 from networkx.drawing.nx_pydot import graphviz_layout
 from io import BytesIO
+from PIL import Image
+
 
 TITLE = """
 <div align="center">
@@ -163,13 +165,16 @@ def create_family_tree(start_model):
 
     plt.title(f"{start_model}'s Family Tree", fontsize=20)
     
-    # Instead of plt.show(), capture the plot as an image in memory
+    # Capture the plot as an image in memory
     img_buffer = BytesIO()
-    plt.savefig(img_buffer, format='png')
+    plt.savefig(img_buffer, format='png', bbox_inches='tight')
     plt.close()
-    img_buffer.seek(0)  # Rewind the buffer to the beginning
+    img_buffer.seek(0)
+
+    # Open the image using PIL
+    img = Image.open(img_buffer)
 
-    return img_buffer.getvalue()
+    return img
     
 with gr.Blocks() as demo:
     gr.Markdown(TITLE)