bug fix

app.py

@@ -148,27 +148,29 @@ def main():
         preds = predict(model, dataloader, st.session_state.date)
         update_plot(df, preds, axs, fig)
 
-    def update_plot2(ax):
-        x = np.linspace(0, 2*np.pi, 100)
-        y = np.sin(x)
-        ax.plot(x, y, label='New Line')
+    
+    @st.cache(allow_output_mutation=True)
+    def create_plot():
+        fig, ax = plt.subplots()
+        x_initial = np.linspace(0, 2*np.pi, 100)
+        y_initial = np.cos(x_initial)
+        ax.plot(x_initial, y_initial, label='Initial Line')
+        ax.set_xlabel('X-axis')
+        ax.set_ylabel('Y-axis')
+        ax.set_title('Plot')
         ax.legend()
-        st.pyplot()
-
-    # Create an initial plot
-    figs, ax = plt.subplots()
-    x_initial = np.linspace(0, 2*np.pi, 100)
-    y_initial = np.cos(x_initial)
-    ax.plot(x_initial, y_initial, label='Initial Line')
-    ax.set_xlabel('X-axis')
-    ax.set_ylabel('Y-axis')
-    ax.set_title('Plot')
-    ax.legend()
-    st.pyplot(figs)
+        return fig
+    
+    fig = create_plot()
+    st.pyplot(fig)
     
-    # Button to update the plot
     if st.button('Add Line'):
-        update_plot2(ax)
+        ax = fig.axes[0]
+        x_new = np.linspace(0, 2*np.pi, 100)
+        y_new = np.sin(x_new)
+        ax.plot(x_new, y_new, label='New Line')
+        ax.legend()
+        st.pyplot(fig)