Update app.py

app.py

@@ -672,48 +672,108 @@ def make_image_sized_pdf(sources):
         return None
 
 
-
-# --- Sidebar Gallery Update Function (Keep from previous) --------
+# --- Sidebar Gallery Update Function (MODIFIED) --------
 def update_gallery():
-    container = st.session_state['asset_gallery_container']
-    container.empty()
-    with container.container():
-        st.markdown("### Asset Gallery 📸📖")
-        cols = st.columns(2)
-        all_files = get_gallery_files()
-        if not all_files: st.info("No assets found."); return
-        files_to_display = all_files[:st.session_state['gallery_size']]
-        for idx, file in enumerate(files_to_display):
-            with cols[idx % 2]:
-                st.session_state['unique_counter'] += 1; unique_id = st.session_state['unique_counter']; basename = os.path.basename(file); st.caption(basename)
+    st.sidebar.markdown("### Asset Gallery 📸📖")
+
+    all_files = get_gallery_files() # Get currently available files
+
+    if not all_files:
+        st.sidebar.info("No assets (images, PDFs, text files) found yet.")
+        return
+
+    st.sidebar.caption(f"Found {len(all_files)} assets:")
+
+    for idx, file in enumerate(all_files):
+        st.session_state['unique_counter'] += 1
+        unique_id = st.session_state['unique_counter']
+        item_key_base = f"gallery_item_{os.path.basename(file)}_{unique_id}"
+        basename = os.path.basename(file)
+        st.sidebar.markdown(f"**{basename}**") # Display filename clearly
+
+        try:
+            file_ext = os.path.splitext(file)[1].lower()
+            # Display previews
+            if file_ext in ['.png', '.jpg', '.jpeg']:
+                 # Add expander for large galleries
+                 with st.sidebar.expander("Preview", expanded=False):
+                      st.image(Image.open(file), use_container_width=True)
+            elif file_ext == '.pdf':
+                 with st.sidebar.expander("Preview (Page 1)", expanded=False):
+                      doc = fitz.open(file)
+                      if len(doc) > 0:
+                          pix = doc[0].get_pixmap(matrix=fitz.Matrix(0.5, 0.5)) # Smaller preview
+                          img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+                          st.image(img, use_container_width=True)
+                      else:
+                          st.warning("Empty PDF")
+                      doc.close()
+            elif file_ext in ['.md', '.txt']:
+                 with st.sidebar.expander("Preview (Start)", expanded=False):
+                      with open(file, 'r', encoding='utf-8', errors='ignore') as f:
+                          content_preview = f.read(200) # Show first 200 chars
+                      st.code(content_preview + "...", language='markdown' if file_ext == '.md' else 'text')
+
+            # --- Actions for the file (Select, Download, Delete) ---
+            action_cols = st.sidebar.columns(3) # Use columns for buttons
+            with action_cols[0]:
+                 checkbox_key = f"cb_{item_key_base}"
+                 st.session_state['asset_checkboxes'][file] = st.checkbox(
+                     "Select",
+                     value=st.session_state['asset_checkboxes'].get(file, False),
+                     key=checkbox_key
+                 )
+            with action_cols[1]:
+                mime_map = {'.png': 'image/png', '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.pdf': 'application/pdf', '.txt': 'text/plain', '.md': 'text/markdown'}
+                mime_type = mime_map.get(file_ext, "application/octet-stream")
+                # Use button for download to avoid complex HTML link generation issues sometimes
+                dl_key = f"dl_{item_key_base}"
                 try:
-                    file_ext = os.path.splitext(file)[1].lower()
-                    if file_ext in ['.png', '.jpg', '.jpeg']: st.image(Image.open(file), use_container_width=True)
-                    elif file_ext == '.pdf':
-                        doc = fitz.open(file)
-                        if len(doc) > 0: 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, use_container_width=True)
-                        else: st.warning("Empty PDF")
-                        doc.close()
-                    elif file_ext in ['.md', '.txt']:
-                        with open(file, 'r', encoding='utf-8', errors='ignore') as f: content_preview = f.read(200)
-                        st.code(content_preview + "...", language='markdown' if file_ext == '.md' else 'text')
-                    checkbox_key = f"asset_cb_{file}_{unique_id}"; st.session_state['asset_checkboxes'][file] = st.checkbox("Select", value=st.session_state['asset_checkboxes'].get(file, False), key=checkbox_key)
-                    mime_map = {'.png': 'image/png', '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.pdf': 'application/pdf', '.txt': 'text/plain', '.md': 'text/markdown'}; mime_type = mime_map.get(file_ext, "application/octet-stream")
-                    st.markdown(get_download_link(file, mime_type, "📥"), unsafe_allow_html=True)
-                    delete_key = f"delete_btn_{file}_{unique_id}"
-                    if st.button("🗑️", key=delete_key, help=f"Delete {basename}"):
-                        try:
-                            os.remove(file); st.session_state['asset_checkboxes'].pop(file, None)
-                            if file in st.session_state.get('layout_snapshots', []): st.session_state['layout_snapshots'].remove(file)
-                            logger.info(f"Deleted asset: {file}"); st.success(f"Deleted {basename}"); st.rerun()
-                        except OSError as e: logger.error(f"Error deleting file {file}: {e}"); st.error(f"Could not delete {basename}")
-                except (fitz.fitz.FileNotFoundError, FileNotFoundError): st.error(f"File not found: {basename}"); st.session_state['asset_checkboxes'].pop(file, None); # Clean up state
-                except (fitz.fitz.FileDataError, fitz.fitz.RuntimeException) as pdf_err: st.error(f"Corrupt PDF: {basename}"); logger.warning(f"Error opening PDF {file}: {pdf_err}")
-                except UnidentifiedImageError: st.error(f"Invalid Image: {basename}"); logger.warning(f"Cannot identify image file {file}")
-                except Exception as e: st.error(f"Error: {basename}"); logger.error(f"Error displaying asset {file}: {e}")
-                st.markdown("---")
-        if len(all_files) > st.session_state['gallery_size']: st.caption(f"Showing {st.session_state['gallery_size']} of {len(all_files)} assets.")
+                    with open(file, "rb") as fp:
+                        st.download_button(
+                             label="📥",
+                             data=fp,
+                             file_name=basename,
+                             mime=mime_type,
+                             key=dl_key,
+                             help="Download this file"
+                         )
+                except Exception as dl_e:
+                     st.error(f"DL Err: {dl_e}")
+
+            with action_cols[2]:
+                delete_key = f"del_{item_key_base}"
+                if st.button("🗑️", key=delete_key, help=f"Delete {basename}"):
+                    try:
+                        os.remove(file)
+                        st.session_state['asset_checkboxes'].pop(file, None) # Remove from selection state
+                        # Remove from layout_snapshots if present
+                        if file in st.session_state.get('layout_snapshots', []):
+                            st.session_state['layout_snapshots'].remove(file)
+                        logger.info(f"Deleted asset: {file}")
+                        st.toast(f"Deleted {basename}!", icon="✅") # Use toast for less intrusive feedback
+                        # REMOVED st.rerun() - Rely on file watcher
+                    except OSError as e:
+                        logger.error(f"Error deleting file {file}: {e}")
+                        st.error(f"Could not delete {basename}")
+                    # Trigger a rerun MANUALLY after deletion completes if file watcher is unreliable
+                    st.rerun()
+
+
+        except (fitz.fitz.FileNotFoundError, FileNotFoundError):
+             st.sidebar.error(f"File not found: {basename}")
+             st.session_state['asset_checkboxes'].pop(file, None) # Clean up state
+        except (fitz.fitz.FileDataError, fitz.fitz.RuntimeException) as pdf_err:
+             st.sidebar.error(f"Corrupt PDF: {basename}")
+             logger.warning(f"Error opening PDF {file}: {pdf_err}")
+        except UnidentifiedImageError:
+            st.sidebar.error(f"Invalid Image: {basename}")
+            logger.warning(f"Cannot identify image file {file}")
+        except Exception as e:
+            st.sidebar.error(f"Error: {basename}")
+            logger.error(f"Error displaying asset {file}: {e}")
 
+        st.sidebar.markdown("---") # Separator between items
 
 # --- UI Elements -----------------------------------------
 
@@ -881,7 +941,7 @@ with tabs[0]:
                 st.image(Image.open(scan_name), caption=f"Scanned: {scan_name}", use_container_width=True)
                 if scan_name not in st.session_state['layout_snapshots']: st.session_state['layout_snapshots'].append(scan_name)
                 st.success(f"Scan saved as {scan_name}")
-                update_gallery(); st.rerun() # Add to gallery
+                update_gallery();  # Add to gallery
             except Exception as e: st.error(f"Failed to save scan: {e}"); logger.error(f"Failed to save camera scan {scan_name}: {e}")
 
         layout_uploads = st.file_uploader("📂 Upload PNG/JPG Images for Layout PDF", type=["png","jpg","jpeg"], accept_multiple_files=True, key="layout_uploader")
@@ -959,7 +1019,7 @@ with tabs[1]:
             try:
                 with open(filename, "wb") as f: f.write(cam0_img.getvalue())
                 st.session_state['cam0_file'] = filename; st.session_state['history'].append(f"Snapshot from Cam 0: {filename}"); st.image(Image.open(filename), caption="Camera 0 Snap", use_container_width=True); logger.info(f"Saved snapshot from Camera 0: {filename}"); st.success(f"Saved {filename}")
-                update_gallery(); st.rerun()
+                update_gallery();
             except Exception as e: 
                 st.error(f"Failed to save Cam 0 snap: {e}"); logger.error(f"Failed to save Cam 0 snap {filename}: {e}")
     with cols[1]:
@@ -974,7 +1034,7 @@ with tabs[1]:
             try:
                 with open(filename, "wb") as f: f.write(cam1_img.getvalue())
                 st.session_state['cam1_file'] = filename; st.session_state['history'].append(f"Snapshot from Cam 1: {filename}"); st.image(Image.open(filename), caption="Camera 1 Snap", use_container_width=True); logger.info(f"Saved snapshot from Camera 1: {filename}"); st.success(f"Saved {filename}")
-                update_gallery(); st.rerun()
+                update_gallery();
             except Exception as e: st.error(f"Failed to save Cam 1 snap: {e}"); logger.error(f"Failed to save Cam 1 snap {filename}: {e}")
 
 
@@ -999,7 +1059,7 @@ with tabs[2]:
                     if download_pdf(url, output_path): st.session_state['downloaded_pdfs'][url] = output_path; logger.info(f"Downloaded PDF from {url} to {output_path}"); st.session_state['history'].append(f"Downloaded PDF: {output_path}"); st.session_state['asset_checkboxes'][output_path] = False; download_count += 1; existing_pdfs.append(output_path)
                     else: st.error(f"Failed to download: {url}")
             status_text.success(f"Download process complete! Successfully downloaded {download_count} new PDFs.")
-            if download_count > 0: update_gallery(); st.rerun()
+            if download_count > 0: update_gallery();
 
     st.subheader("Create Snapshots from Gallery PDFs")
     snapshot_mode = st.selectbox("Snapshot Mode", ["First Page (High-Res)", "First Two Pages (High-Res)", "All Pages (High-Res)", "First Page (Low-Res Preview)"], key="pdf_snapshot_mode")
@@ -1019,7 +1079,7 @@ with tabs[2]:
                     st.write(f"Snapshots for {os.path.basename(pdf_path)}:"); cols = st.columns(3)
                     for i, snap_path in enumerate(new_snapshots):
                          with cols[i % 3]: st.image(Image.open(snap_path), caption=os.path.basename(snap_path), use_container_width=True); st.session_state['asset_checkboxes'][snap_path] = False # Add to gallery
-            if total_snapshots_generated > 0: st.success(f"Generated {total_snapshots_generated} snapshots from {snapshot_count} PDFs."); update_gallery(); st.rerun()
+            if total_snapshots_generated > 0: st.success(f"Generated {total_snapshots_generated} snapshots from {snapshot_count} PDFs."); update_gallery();
             else: st.warning("No snapshots were generated. Check logs or PDF files.")
 
 # --- Tab 4: Build Titan (Local Models) ---
@@ -1162,8 +1222,6 @@ with tabs[3]:
                 finally:
                      progress_bar_build.progress(1.0)
                      status_text_build.empty()
-                     # Update sidebar selector
-                     st.rerun()
 
         st.subheader("Manage Local Models")
         loaded_model_paths = list(st.session_state.get('local_models', {}).keys())
@@ -1196,7 +1254,6 @@ with tabs[3]:
                             shutil.rmtree(path_to_delete)
                         st.success(f"Deleted model '{model_to_delete}' and its files.")
                         logger.info(f"Deleted local model: {path_to_delete}")
-                        st.rerun()
                     except Exception as e:
                         st.error(f"Failed to delete model '{model_to_delete}': {e}")
                         logger.error(f"Failed to delete model {path_to_delete}: {e}")