Update app.py

app.py

@@ -9,11 +9,10 @@ import re
 import random
 import torch
 import time
-import time
+import shutil  # Added for zip functionality
 
 from PIL import Image
 from io import BytesIO
-from PIL import Image
 from diffusers import DiffusionPipeline, LCMScheduler, AutoencoderTiny
 
 try:
@@ -33,6 +32,77 @@ device = torch.device(
 torch_device = device
 torch_dtype = torch.float16
 
+
+# add file save and download and clear:
+# 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"
+    with zipfile.ZipFile(zip_filename, 'w') as zipf:
+        for file in files:
+            zipf.write(file, os.path.basename(file))
+    return zip_filename
+
+# Function to encode a file to base64
+def encode_file_to_base64(file_path):
+    with open(file_path, "rb") as file:
+        encoded = base64.b64encode(file.read()).decode()
+    return encoded
+
+
+# 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
+    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 clear all image files
+def clear_all_images():
+    base_dir = os.getcwd()  # Get the current base directory
+    img_files = [file for file in os.listdir(base_dir) if file.lower().endswith((".png", ".jpg", ".jpeg"))]  # List all files ending with ".jpg" or ".jpeg"
+    
+    # Remove all image files
+    for file in img_files:
+        os.remove(file)
+
+
+# Add "Save All" button with emoji
+save_all_button = gr.Button("💾 Save All", scale=1)
+# Add "Clear All" button with emoji
+clear_all_button = gr.Button("🗑️ Clear All", scale=1)
+
+# 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 download_link:
+        gr.write(download_link)
+
+# Function to handle "Clear All" button click
+def clear_all_button_click():
+    clear_all_images()
+
+# Attach click event handlers to the buttons
+save_all_button.click(save_all_button_click)
+clear_all_button.click(clear_all_button_click)
+
+# Add buttons to the Streamlit app
+gr.button(save_all_button)
+gr.button(clear_all_button)
+
+
+
+
+
+
+
 print(f"SAFETY_CHECKER: {SAFETY_CHECKER}")
 print(f"TORCH_COMPILE: {TORCH_COMPILE}")
 print(f"device: {device}")