Update app.py

app.py

@@ -80,7 +80,7 @@ if uploaded_file is not None:
     # Add a slider to change the scale
     with col1:
         st.markdown(
-            """<style>[data-baseweb="select"] {margin-top: -50px;}</style>""",
+            """<style>[data-baseweb="select"] {margin-top: -52px;}</style>""",
             unsafe_allow_html=True
         )
     
@@ -110,12 +110,12 @@ if uploaded_file is not None:
 
         plot_prediction(y_pred)
 
-    # with col2:
-    #     download = st.button('Download FITS File')
-
-    # if download:
-    #     ccd = CCDData(pred, unit="adu", wcs=wcs)
-    #     ccd.write(f"predicted.fits", overwrite=True)
-    #     with open('predicted.fits', 'rb') as f:
-    #         data = f.read()
-    #     st.download_button(label="Download", data=data, file_name="predicted.fits", mime="application/octet-stream")
+    with col2:
+        download = st.button('Download')
+
+    if download:
+        ccd = CCDData(pred, unit="adu", wcs=wcs)
+        ccd.write("predicted.fits", overwrite=True)
+        with open('predicted.fits', 'rb') as f:
+            data = f.read()
+        st.download_button(label="Download", data=data, file_name="predicted.fits", mime="application/octet-stream")