Update app.py

app.py

@@ -3,8 +3,6 @@ import streamlit as st
 import requests
 import re
 from dotenv import load_dotenv
-from PIL import Image
-from io import BytesIO
 import tempfile
 
 # Load environment variables (for FLOWISE_API token)
@@ -51,11 +49,10 @@ if st.button("Generate Diagram"):
                 with open(temp_file, "w") as f:
                     f.write(python_code)
 
-                # Execute the code
+                # Execute the code and look for 'diagram.jpg'
                 try:
                     exec(open(temp_file).read())
-                    # Assuming the diagram was saved as "WordPress_Architecture.png" in the script
-                    diagram_path = os.path.join(tmpdirname, "WordPress_Architecture.png")
+                    diagram_path = os.path.join(tmpdirname, "diagram.jpg")
 
                     # Display the generated image
                     if os.path.exists(diagram_path):