Update app.py

app.py

@@ -1,7 +1,6 @@
 import streamlit as st
 from PIL import Image, ImageEnhance
 from io import BytesIO
-import numpy as np
 import zipfile
 import os
 from fpdf import FPDF
@@ -52,7 +51,6 @@ if uploaded_files:
         st.markdown(f"---\n### 🖼️ {file.name}")
         original = Image.open(BytesIO(file_bytes)).convert("RGB")
 
-        # Init state
         if f"original_{key}" not in st.session_state:
             st.session_state[f"original_{key}"] = original
             st.session_state[f"angle_{key}"] = 0
@@ -113,8 +111,10 @@ if uploaded_files:
 
     # --- EXPORT SECTION ---
     st.markdown("### 🧾 Export Options")
-    col1, col2 = st.columns([1, 1])
-    with col1:
+
+    # ZIP inline layout
+    zip_col1, zip_col2 = st.columns([1, 2])
+    with zip_col1:
         if st.button("📦 Generate ZIP"):
             zip_buffer = BytesIO()
             with zipfile.ZipFile(zip_buffer, "w") as zipf:
@@ -123,9 +123,14 @@ if uploaded_files:
                     img.save(buffer, format="PNG")
                     zipf.writestr(f"{name}.png", buffer.getvalue())
             zip_buffer.seek(0)
-            st.download_button("⬇️ Download ZIP", zip_buffer, "images.zip", mime="application/zip", key="zip_dl")
-
-    with col2:
+            st.session_state["zip_file"] = zip_buffer
+    with zip_col2:
+        if "zip_file" in st.session_state:
+            st.download_button("⬇️ Download ZIP", st.session_state["zip_file"], "images.zip", mime="application/zip", key="zip_dl")
+
+    # PDF inline layout
+    pdf_col1, pdf_col2 = st.columns([1, 2])
+    with pdf_col1:
         if st.button("📄 Generate PDF"):
             pdf = FPDF()
             tmp_paths = []
@@ -138,6 +143,9 @@ if uploaded_files:
             out_path = "/tmp/final.pdf"
             pdf.output(out_path)
             with open(out_path, "rb") as f:
-                st.download_button("⬇️ Download PDF", f, "scanned.pdf", mime="application/pdf", key="pdf_dl")
+                st.session_state["pdf_file"] = f.read()
             for p in tmp_paths + [out_path]:
                 os.remove(p)
+    with pdf_col2:
+        if "pdf_file" in st.session_state:
+            st.download_button("⬇️ Download PDF", st.session_state["pdf_file"], "scanned.pdf", mime="application/pdf", key="pdf_dl")