Update app.py

app.py

@@ -178,16 +178,20 @@ if uploaded_file is not None:
         threshold = st.slider("Threshold", 0.0, 1.0, 0.0, 0.05) #, label_visibility="hidden")
         
     # Decompose button
-    with col5: decompose = st.button('Decompose')
+    with col5: decompose = st.button('Decompose', key="decompose")
 
     with col6:
-        with open("blank.txt", 'w') as f: pass
-        with open("blank.txt", 'rb') as f: res = f.read()
         st.markdown("")
         st.markdown("")
         # st.markdown("""<style>[data-baseweb="select"] {margin-top: 16px;}</style>""", unsafe_allow_html=True)
-        fname = uploaded_file.name.strip(".fits")    
-        download = st.download_button(label="Download", data=res, file_name=f'{fname}_{int(scale*128)}.zip', disabled=True, mime="application/octet-stream")
+        fname = uploaded_file.name.strip(".fits")
+        shutil.make_archive("predictions", 'zip', "predictions")
+        with open('predictions.zip', 'rb') as f:
+            res = f.read()
+        download = st.download_button(label="Download", data=res, 
+                                      file_name=f'{fname}_{int(scale*128)}.zip', 
+                                      disabled=st.session_state.get("disabled", True), 
+                                      mime="application/octet-stream")
         
     # Make two columns for plots
     _, colA, colB, colC, _ = st.columns([bordersize,1,1,1,bordersize])
@@ -203,14 +207,10 @@ if uploaded_file is not None:
         plot_prediction(y_pred_th)
 
         print(type(download))
-        if decompose or download:
-            download.disabled = False
-            
+        if decompose or download:            
             image_decomposed = decompose_cavity(y_pred_th)
     
             plot_decomposed(image_decomposed)
 
-            shutil.make_archive("predictions", 'zip', "predictions")
-            with open('predictions.zip', 'rb') as f:
-                res = f.read()
-            download.data=res
+        if st.session_state.get("decompose", True):
+            st.session_state.disabled = False