Update app.py

app.py

@@ -10,10 +10,11 @@ import random
 import torch
 import time
 import shutil  # Added for zip functionality
-import zipfile
 from PIL import Image
 from io import BytesIO
 from diffusers import DiffusionPipeline, LCMScheduler, AutoencoderTiny
+import zipfile
+import glob
 
 try:
     import intel_extension_for_pytorch as ipex
@@ -33,6 +34,44 @@ torch_device = device
 torch_dtype = torch.float16
 
 
+# Function to create a zip file from a list of files
+def create_zip(files):
+    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+    zip_filename = f"images_{timestamp}.zip"
+    print('Creating file ' + zip_filename)
+    with zipfile.ZipFile(zip_filename, 'w') as zipf:
+        for file in files:
+            zipf.write(file, os.path.basename(file))
+            print('added:' + file)
+    return zip_filename
+    
+# Function to save all images as a zip file and provide a base64 download link
+def save_all_images(images):
+    if len(images) == 0:
+        return None, None
+    zip_filename = create_zip(images)  # Create a zip file from the list of image files
+    print(' Zip file created:' + zip_filename)
+
+    zip_base64 = encode_file_to_base64(zip_filename)  # Encode the zip file to base64
+    download_link = f'<a href="data:application/zip;base64,{zip_base64}" download="{zip_filename}">Download All</a>'
+    
+    return zip_filename, download_link
+        
+# Function to handle "Save All" button click
+def save_all_button_click():
+    images = [file for file in os.listdir() if file.lower().endswith((".png", ".jpg", ".jpeg"))]
+    zip_filename, download_link = save_all_images(images)   
+    if zip_filename:
+        print(zip_filename)
+        gr.Button(link=zip_filename)
+    if download_link:
+        print(download_link)
+        #gr.HTML(download_link)
+        gr.Button(link=download_link)
+
+# Function to handle "Clear All" button click
+def clear_all_button_click():
+    clear_all_images()
 
 # Function to encode a file to base64
 def encode_file_to_base64(file_path):