Update app.py

app.py

@@ -82,7 +82,7 @@ def arrange_images(image_files, canvas_size=(3000, 3000)):
     buffer = BytesIO()
     canvas.save(buffer, format="PNG")
     buffer.seek(0)
-    return buffer
+    return buffer, canvas
 
 # Streamlit App
 st.title("Dynamic Dungeon Map Generator")
@@ -95,6 +95,17 @@ map_dir = "."  # Top-level directory
 canvas_size = 3000
 scroll_offset = {"x": 0, "y": 0}
 
+# Generate map on load
+if "layout_image" not in st.session_state:
+    # Scan the directory for .png files
+    image_files = [f for f in os.listdir(map_dir) if f.endswith(".png")]
+    if image_files:
+        layout_image, canvas = arrange_images(image_files, canvas_size=(canvas_size, canvas_size))
+        st.session_state["layout_image"] = layout_image
+        st.session_state["canvas"] = canvas
+    else:
+        st.write("No PNG files found in the top-level directory.")
+
 # WASD Controls for scrolling
 st.sidebar.title("Scroll Map")
 scroll_buttons = st.sidebar.columns(3)
@@ -107,26 +118,32 @@ if scroll_buttons[2].button("➡️"):
 if st.sidebar.button("⬇️"):
     scroll_offset["y"] += 100
 
-# Scan the directory for .png files
-image_files = [f for f in os.listdir(map_dir) if f.endswith(".png")]
-
-if image_files:
-    # Add "Create Map" button
-    if st.button("🗺️ Create Map"):
-        scroll_offset = {"x": 0, "y": 0}  # Reset scroll offset
-        layout_image = arrange_images(image_files, canvas_size=(canvas_size, canvas_size))
-        if layout_image:
-            st.image(
-                layout_image,
-                caption="Generated Dungeon Map Layout",
-                use_container_width=True,
-                output_format="PNG",
-                clamp=True
-            )
-        else:
-            st.write("Failed to generate a map. Please check the images in the directory.")
-else:
-    st.write("No PNG files found in the top-level directory.")
+# Display map
+if "layout_image" in st.session_state:
+    st.image(
+        st.session_state["layout_image"],
+        caption="Generated Dungeon Map Layout",
+        use_column_width=True,
+        output_format="PNG",
+        clamp=True,
+    )
+
+# Add save button
+if st.button("💾 Save Map"):
+    save_path = "saved_dungeon_map.png"
+    st.session_state["canvas"].save(save_path)
+    st.success(f"Map saved as {save_path}")
+
+# Regenerate map button
+if st.button("🗺️ Regenerate Map"):
+    image_files = [f for f in os.listdir(map_dir) if f.endswith(".png")]
+    if image_files:
+        layout_image, canvas = arrange_images(image_files, canvas_size=(canvas_size, canvas_size))
+        st.session_state["layout_image"] = layout_image
+        st.session_state["canvas"] = canvas
+        st.experimental_rerun()
+    else:
+        st.write("No PNG files found in the top-level directory.")
 
 # Sidebar for file upload
 st.sidebar.title("Options")