Update app.py

app.py

@@ -330,7 +330,7 @@ def get_model_files(model_type="causal_lm"):
     return [d for d in glob.glob(path) if os.path.isdir(d)]
 
 def get_gallery_files(file_types=["png", "pdf"]):
-    return sorted([f for ext in file_types for f in glob.glob(f"*.{ext}")])
+    return sorted(list(set([f for ext in file_types for f in glob.glob(f"*.{ext}")])))  # Deduplicate files
 
 def get_pdf_files():
     return sorted(glob.glob("*.pdf"))
@@ -539,17 +539,18 @@ def update_gallery():
                     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(
+                checkbox_key = f"asset_{file}_{idx}"  # Unique key with index
+                st.session_state['asset_checkboxes'][file] = st.checkbox(
                     "Use for SFT/Input", 
-                    value=st.session_state['asset_checkboxes'].get(checkbox_key, False), 
+                    value=st.session_state['asset_checkboxes'].get(file, 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}"):
+                if st.button("Zap It! 🗑️", key=f"delete_{file}_{idx}"):  # Unique key with index
                     os.remove(file)
-                    del st.session_state['asset_checkboxes'][checkbox_key]
+                    if file in st.session_state['asset_checkboxes']:
+                        del st.session_state['asset_checkboxes'][file]
                     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:
@@ -664,7 +665,7 @@ with tab2:
 
     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 path in get_gallery_files() if path.endswith('.pdf') and st.session_state['asset_checkboxes'].get(f"asset_{path}", False)]
+        selected_pdfs = [path for path in get_gallery_files() if path.endswith('.pdf') and st.session_state['asset_checkboxes'].get(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 (High-Res)": "allpages"}[mode]
@@ -672,7 +673,7 @@ with tab2:
                 for snapshot in snapshots:
                     st.image(Image.open(snapshot), caption=snapshot, use_container_width=True)
         else:
-            st.warning("No PDFs selected for snapshotting! Check some boxes first. 📝")
+            st.warning("No PDFs selected for snapshotting! Check some boxes in the sidebar gallery.")
 
 with tab3:
     st.header("Build Titan 🌱")
@@ -732,7 +733,7 @@ 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):
-            selected_files = [path for path in get_gallery_files() if st.session_state['asset_checkboxes'].get(f"asset_{path}", False)]
+            selected_files = [path for path in get_gallery_files() if st.session_state['asset_checkboxes'].get(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")
@@ -829,7 +830,7 @@ with tab6:
 
 with tab7:
     st.header("Test OCR 🔍")
-    all_files = [path for path in get_gallery_files() if st.session_state['asset_checkboxes'].get(f"asset_{path}", False)]
+    all_files = [path for path in get_gallery_files() if st.session_state['asset_checkboxes'].get(path, False)]
     if all_files:
         selected_file = st.selectbox("Select Image or PDF", all_files, key="ocr_select")
         if selected_file:
@@ -856,7 +857,7 @@ with tab7:
 
 with tab8:
     st.header("Test Image Gen 🎨")
-    all_files = [path for path in get_gallery_files() if st.session_state['asset_checkboxes'].get(f"asset_{path}", False)]
+    all_files = [path for path in get_gallery_files() if st.session_state['asset_checkboxes'].get(path, False)]
     if all_files:
         selected_file = st.selectbox("Select Image or PDF", all_files, key="gen_select")
         if selected_file:
@@ -885,7 +886,7 @@ with tab8:
 with tab9:
     st.header("Custom Diffusion 🎨🤓")
     st.write("Unleash your inner artist with our tiny diffusion models!")
-    all_files = [path for path in get_gallery_files() if st.session_state['asset_checkboxes'].get(f"asset_{path}", False)]
+    all_files = [path for path in get_gallery_files() if st.session_state['asset_checkboxes'].get(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")