Enhance Gallery UI: add clear, download zip, improve layout

ui/layout.py

@@ -1,4 +1,7 @@
 import os
+import shutil
+import zipfile
+import tempfile
 import gradio as gr
 from core.settings import *
 
@@ -35,6 +38,36 @@ def build_ui(event_handler_function):
                     image_paths.append(os.path.join(output_dir, fname))
         return image_paths
 
+    # Clear all images from the gallery folder and return an empty list for the UI.
+    def clear_gallery_images():
+        """Remove every file in ``OUTPUT_DIR`` and return an empty list for the gallery.
+        This provides a quick way for users to reset the gallery view.
+        """
+        output_dir = os.path.abspath(OUTPUT_DIR)
+        if os.path.isdir(output_dir):
+            for f in os.listdir(output_dir):
+                try:
+                    os.remove(os.path.join(output_dir, f))
+                except Exception:
+                    pass
+        return []
+
+    # Create a zip archive of all images in the gallery for downloading.
+    def download_gallery_zip():
+        """Package all files in ``OUTPUT_DIR`` into a temporary zip file.
+        Returns the file path which Gradio will serve as a downloadable file.
+        """
+        output_dir = os.path.abspath(OUTPUT_DIR)
+        tmp_fd, tmp_path = tempfile.mkstemp(suffix=".zip")
+        os.close(tmp_fd)
+        with zipfile.ZipFile(tmp_path, "w", zipfile.ZIP_DEFLATED) as zf:
+            if os.path.isdir(output_dir):
+                for fname in sorted(os.listdir(output_dir)):
+                    file_path = os.path.join(output_dir, fname)
+                    if os.path.isfile(file_path):
+                        zf.write(file_path, arcname=fname)
+        return tmp_path
+
     # All UI components, including tabs, must be created inside the same Gradio Blocks context.
     with gr.Blocks() as demo:
         gr.Markdown("# ImageGen")
@@ -69,12 +102,20 @@ def build_ui(event_handler_function):
 
             # New Gallery tab to display generated images
             with gr.TabItem("Gallery", id=5):
-                # Refresh button to manually reload thumbnails from the output folder.
-                refresh_btn = gr.Button("Refresh Gallery", variant="secondary")
-                gallery = gr.Gallery(label="Generated Images", show_label=False, columns=3, height="auto", type="filepath")
+                # Row of control buttons for the gallery.
+                with gr.Row():
+                    refresh_btn = gr.Button("Refresh Gallery", variant="secondary")
+                    clear_btn = gr.Button("Clear Gallery", variant="secondary")
+                    download_btn = gr.Button("Download All", variant="secondary")
+                # Main gallery component.
+                gallery = gr.Gallery(label="Generated Images", show_label=False, columns=4, height="auto", type="filepath")
                 ui_components["gallery"] = gallery
-                # Bind the refresh button to the same loader used on startup.
+                # Hidden file component for zip download.
+                download_file = gr.File(label="Download", visible=False)
+                # Bind buttons.
                 refresh_btn.click(fn=load_gallery_images, outputs=gallery)
+                clear_btn.click(fn=clear_gallery_images, outputs=gallery)
+                download_btn.click(fn=download_gallery_zip, outputs=download_file)
 
             ui_components["tabs"] = tabs
             ui_components["image_gen_tabs"] = tabs