Update app.py

app.py

@@ -59,8 +59,8 @@ if 'model_loaded' not in st.session_state:
     st.session_state['model_loaded'] = False
 if 'processing' not in st.session_state:
     st.session_state['processing'] = {}
-if 'pdf_checkboxes' not in st.session_state:
-    st.session_state['pdf_checkboxes'] = {}
+if 'asset_checkboxes' not in st.session_state:
+    st.session_state['asset_checkboxes'] = {}
 if 'downloaded_pdfs' not in st.session_state:
     st.session_state['downloaded_pdfs'] = {}
 
@@ -329,7 +329,7 @@ def get_model_files(model_type="causal_lm"):
     path = "models/*" if model_type == "causal_lm" else "diffusion_models/*"
     return [d for d in glob.glob(path) if os.path.isdir(d)]
 
-def get_gallery_files(file_types=["png"]):
+def get_gallery_files(file_types=["png", "pdf"]):
     return sorted([f for ext in file_types for f in glob.glob(f"*.{ext}")])
 
 def get_pdf_files():
@@ -367,11 +367,11 @@ async def process_pdf_snapshot(pdf_path, mode="single"):
                 output_file = generate_filename(f"twopage_{i}", "png")
                 pix.save(output_file)
                 output_files.append(output_file)
-        elif mode == "allthumbs":
+        elif mode == "allpages":
             for i in range(len(doc)):
                 page = doc[i]
-                pix = page.get_pixmap(matrix=fitz.Matrix(0.5, 0.5))
-                output_file = generate_filename(f"thumb_{i}", "png")
+                pix = page.get_pixmap(matrix=fitz.Matrix(2.0, 2.0))
+                output_file = generate_filename(f"page_{i}", "png")
                 pix.save(output_file)
                 output_files.append(output_file)
         doc.close()
@@ -506,19 +506,56 @@ def calculate_cargo_travel_time(origin_coords: Tuple[float, float], destination_
 st.title("AI Vision & SFT Titans 🚀")
 
 st.sidebar.header("Captured Files 📜")
+cols = st.sidebar.columns(2)
+with cols[0]:
+    if st.button("Zip All 🤐"):
+        zip_path = f"all_assets_{int(time.time())}.zip"
+        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            for file in get_gallery_files():
+                zipf.write(file, os.path.basename(file))
+        st.sidebar.markdown(get_download_link(zip_path, "application/zip", "Download All Assets"), unsafe_allow_html=True)
+with cols[1]:
+    if st.button("Zap All! 🗑️"):
+        for file in get_gallery_files():
+            os.remove(file)
+        st.session_state['asset_checkboxes'].clear()
+        st.session_state['downloaded_pdfs'].clear()
+        st.sidebar.success("All assets vaporized! 💨")
+        st.rerun()
+
 gallery_size = st.sidebar.slider("Gallery Size", 1, 10, 2)
 def update_gallery():
-    media_files = get_gallery_files(["png"])
-    pdf_files = get_pdf_files()
-    if media_files or pdf_files:
-        st.sidebar.subheader("Images 📸")
+    all_files = get_gallery_files()
+    if all_files:
+        st.sidebar.subheader("Asset Gallery 📸📖")
         cols = st.sidebar.columns(2)
-        for idx, file in enumerate(media_files[:gallery_size * 2]):
+        for idx, file in enumerate(all_files[:gallery_size * 2]):
             with cols[idx % 2]:
-                st.image(Image.open(file), caption=os.path.basename(file), use_container_width=True)
-        st.sidebar.subheader("PDF Downloads 📖")
-        for pdf_file in pdf_files[:gallery_size * 2]:
-            st.markdown(get_download_link(pdf_file, "application/pdf", f"📥 Grab {os.path.basename(pdf_file)}"), unsafe_allow_html=True)
+                if file.endswith('.png'):
+                    st.image(Image.open(file), caption=os.path.basename(file), use_container_width=True)
+                else:
+                    doc = fitz.open(file)
+                    pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5))
+                    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+                    st.image(img, caption=os.path.basename(file), use_container_width=True)
+                    doc.close()
+                checkbox_key = f"asset_{file}"
+                st.session_state['asset_checkboxes'][checkbox_key] = st.checkbox(
+                    "Use for SFT/Input", 
+                    value=st.session_state['asset_checkboxes'].get(checkbox_key, False), 
+                    key=checkbox_key
+                )
+                mime_type = "image/png" if file.endswith('.png') else "application/pdf"
+                st.markdown(get_download_link(file, mime_type, "Snag It! 📥"), unsafe_allow_html=True)
+                if st.button("Zap It! 🗑️", key=f"delete_{file}"):
+                    os.remove(file)
+                    del st.session_state['asset_checkboxes'][checkbox_key]
+                    if file.endswith('.pdf'):
+                        url_key = next((k for k, v in st.session_state['downloaded_pdfs'].items() if v == file), None)
+                        if url_key:
+                            del st.session_state['downloaded_pdfs'][url_key]
+                    st.sidebar.success(f"Asset {os.path.basename(file)} vaporized! 💨")
+                    st.rerun()
 update_gallery()
 
 st.sidebar.subheader("Model Management 🗂️")
@@ -625,44 +662,12 @@ with tab2:
         status_text.text("Robo-Download complete! 🚀")
         update_gallery()
 
-    st.subheader("PDF Gallery 📖")
-    downloaded_pdfs = list(st.session_state['downloaded_pdfs'].values())
-    if downloaded_pdfs:
-        cols_per_row = 3
-        for i in range(0, len(downloaded_pdfs), cols_per_row):
-            cols = st.columns(cols_per_row)
-            for j, pdf_path in enumerate(downloaded_pdfs[i:i + cols_per_row]):
-                with cols[j]:
-                    doc = fitz.open(pdf_path)
-                    page = doc[0]
-                    pix = page.get_pixmap(matrix=fitz.Matrix(0.5, 0.5))
-                    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-                    st.image(img, caption=os.path.basename(pdf_path), use_container_width=True)
-                    checkbox_key = f"pdf_{pdf_path}"
-                    st.session_state['pdf_checkboxes'][checkbox_key] = st.checkbox(
-                        "Use for SFT/Input", 
-                        value=st.session_state['pdf_checkboxes'].get(checkbox_key, False), 
-                        key=checkbox_key
-                    )
-                    st.markdown(get_download_link(pdf_path, "application/pdf", "Snag It! 📥"), unsafe_allow_html=True)
-                    if st.button("Zap It! 🗑️", key=f"delete_{pdf_path}"):
-                        os.remove(pdf_path)
-                        url_key = next((k for k, v in st.session_state['downloaded_pdfs'].items() if v == pdf_path), None)
-                        if url_key:
-                            del st.session_state['downloaded_pdfs'][url_key]
-                        del st.session_state['pdf_checkboxes'][checkbox_key]
-                        st.success(f"PDF {os.path.basename(pdf_path)} vaporized! 💨")
-                        st.rerun()
-                    doc.close()
-    else:
-        st.info("No PDFs captured yet. Feed the robo-downloader some URLs! 🤖")
-
-    mode = st.selectbox("Snapshot Mode", ["Single Page (High-Res)", "Two Pages (High-Res)", "All Pages (Thumbnails)"], key="download_mode")
+    mode = st.selectbox("Snapshot Mode", ["Single Page (High-Res)", "Two Pages (High-Res)", "All Pages (High-Res)"], key="download_mode")
     if st.button("Snapshot Selected 📸"):
-        selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)]
+        selected_pdfs = [path for path in get_gallery_files() if path.endswith('.pdf') and st.session_state['asset_checkboxes'].get(f"asset_{path}", False)]
         if selected_pdfs:
             for pdf_path in selected_pdfs:
-                mode_key = {"Single Page (High-Res)": "single", "Two Pages (High-Res)": "twopage", "All Pages (Thumbnails)": "allthumbs"}[mode]
+                mode_key = {"Single Page (High-Res)": "single", "Two Pages (High-Res)": "twopage", "All Pages (High-Res)": "allpages"}[mode]
                 snapshots = asyncio.run(process_pdf_snapshot(pdf_path, mode_key))
                 for snapshot in snapshots:
                     st.image(Image.open(snapshot), caption=snapshot, use_container_width=True)
@@ -727,12 +732,9 @@ with tab4:
                 st.markdown(get_download_link(zip_path, "application/zip", "Download Fine-Tuned Titan"), unsafe_allow_html=True)
                 st.rerun()
         elif isinstance(st.session_state['builder'], DiffusionBuilder):
-            captured_files = get_gallery_files(["png"])
-            selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)]
-            if len(captured_files) + len(selected_pdfs) >= 2:
-                demo_data = [{"image": img, "text": f"Superhero {os.path.basename(img).split('.')[0]}"} for img in captured_files]
-                for pdf_path in selected_pdfs:
-                    demo_data.append({"image": pdf_path, "text": f"PDF {os.path.basename(pdf_path)}"})
+            selected_files = [path for path in get_gallery_files() if st.session_state['asset_checkboxes'].get(f"asset_{path}", False)]
+            if len(selected_files) >= 2:
+                demo_data = [{"image": file, "text": f"Asset {os.path.basename(file).split('.')[0]}"} for file in selected_files]
                 edited_data = st.data_editor(pd.DataFrame(demo_data), num_rows="dynamic")
                 if st.button("Fine-Tune with Dataset 🔄"):
                     images = [Image.open(row["image"]) if row["image"].endswith('.png') else Image.frombytes("RGB", fitz.open(row["image"])[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)).size, fitz.open(row["image"])[0].get_pixmap(matrix=fitz.Matrix(2.0, 2.0)).samples) for _, row in edited_data.iterrows()]
@@ -785,7 +787,6 @@ with tab5:
                 status_container.empty()
         elif isinstance(st.session_state['builder'], DiffusionBuilder):
             test_prompt = st.text_area("Enter Test Prompt", "Neon Batman")
-            selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)]
             if st.button("Run Test ▶️"):
                 image = st.session_state['builder'].generate(test_prompt)
                 output_file = generate_filename("diffusion_test", "png")
@@ -828,9 +829,7 @@ with tab6:
 
 with tab7:
     st.header("Test OCR 🔍")
-    captured_files = get_gallery_files(["png"])
-    selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)]
-    all_files = captured_files + selected_pdfs
+    all_files = [path for path in get_gallery_files() if st.session_state['asset_checkboxes'].get(f"asset_{path}", False)]
     if all_files:
         selected_file = st.selectbox("Select Image or PDF", all_files, key="ocr_select")
         if selected_file:
@@ -853,13 +852,11 @@ with tab7:
                 st.success(f"OCR output saved to {output_file}")
                 st.session_state['processing']['ocr'] = False
     else:
-        st.warning("No images or PDFs captured yet. Use Camera Snap or Download PDFs first!")
+        st.warning("No images or PDFs selected yet. Check some boxes in the sidebar gallery!")
 
 with tab8:
     st.header("Test Image Gen 🎨")
-    captured_files = get_gallery_files(["png"])
-    selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)]
-    all_files = captured_files + selected_pdfs
+    all_files = [path for path in get_gallery_files() if st.session_state['asset_checkboxes'].get(f"asset_{path}", False)]
     if all_files:
         selected_file = st.selectbox("Select Image or PDF", all_files, key="gen_select")
         if selected_file:
@@ -883,14 +880,12 @@ with tab8:
                 st.success(f"Image saved to {output_file}")
                 st.session_state['processing']['gen'] = False
     else:
-        st.warning("No images or PDFs captured yet. Use Camera Snap or Download PDFs first!")
+        st.warning("No images or PDFs selected yet. Check some boxes in the sidebar gallery!")
 
 with tab9:
     st.header("Custom Diffusion 🎨🤓")
     st.write("Unleash your inner artist with our tiny diffusion models!")
-    captured_files = get_gallery_files(["png"])
-    selected_pdfs = [path for key, path in st.session_state['downloaded_pdfs'].items() if st.session_state['pdf_checkboxes'].get(f"pdf_{path}", False)]
-    all_files = captured_files + selected_pdfs
+    all_files = [path for path in get_gallery_files() if st.session_state['asset_checkboxes'].get(f"asset_{path}", False)]
     if all_files:
         st.subheader("Select Images or PDFs to Train")
         selected_files = st.multiselect("Pick Images or PDFs", all_files, key="diffusion_select")
@@ -929,6 +924,6 @@ with tab9:
             st.success(f"Image saved to {output_file}")
             st.session_state['processing']['diffusion'] = False
     else:
-        st.warning("No images or PDFs captured yet. Use Camera Snap or Download PDFs first!")
+        st.warning("No images or PDFs selected yet. Check some boxes in the sidebar gallery!")
 
 update_gallery()