Update backup.app.py

backup.app.py

@@ -9,11 +9,10 @@ import re
 import random
 import torch
 import time
-import time
-
+import shutil  # Added for zip functionality
+import zipfile
 from PIL import Image
 from io import BytesIO
-from PIL import Image
 from diffusers import DiffusionPipeline, LCMScheduler, AutoencoderTiny
 
 try:
@@ -33,6 +32,97 @@ device = torch.device(
 torch_device = device
 torch_dtype = torch.float16
 
+
+
+# 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
+
+
+    
+def create_zip_of_files(files):
+    """
+    Create a zip file from a list of files.
+    """
+    zip_name = "all_files.zip"
+    with zipfile.ZipFile(zip_name, 'w') as zipf:
+        for file in files:
+            zipf.write(file)
+    return zip_name
+
+
+def get_zip_download_link(zip_file):
+    """
+    Generate a link to download the zip file.
+    """
+    with open(zip_file, 'rb') as f:
+        data = f.read()
+    b64 = base64.b64encode(data).decode()
+    href = f'<a href="data:application/zip;base64,{b64}" download="{zip_file}">Download All</a>'
+    return href
+
+    
+
+# 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)
+        print('removed:' + file)
+
+        
+# 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"
+    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)
+
+    # gr.Button(link="/file=" + zip_filename)
+
+    # remove?
+    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>'
+
+    # redirect_button = gr.Button("Clear", variant='secondary')
+    # redirect_button.click(None, None,None, _js="window.location.assign('https://google.com');")
+    
+    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)
+        # gr.File(value=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()
+
 print(f"SAFETY_CHECKER: {SAFETY_CHECKER}")
 print(f"TORCH_COMPILE: {TORCH_COMPILE}")
 print(f"device: {device}")
@@ -105,23 +195,25 @@ def predict(prompt, guidance, steps, seed=1231231):
     if nsfw_content_detected:
         nsfw=gr.Button("🕹️NSFW🎨", scale=1)
 
-    central = pytz.timezone('US/Central')
-    safe_date_time = datetime.datetime.now().strftime("%Y%m%d")
-    replaced_prompt = prompt.replace(" ", "_").replace("\n", "_")
-    safe_prompt = "".join(x for x in replaced_prompt if x.isalnum() or x == "_")[:90]
-    filename = f"{safe_date_time}_{safe_prompt}.png"
-    
-
-    # Save the image
-    if len(results.images) > 0:
-        image_path = os.path.join("", filename)  # Specify your directory
-        results.images[0].save(image_path)
-        print(f"#Image saved as {image_path}")
-        encoded_image = encode_image(image)
-        html_link = f'<a href="data:image/png;base64,{encoded_image}" download="{filename}">Download Image</a>'
-        gr.Markdown(html_link)
-    
-
+    try: 
+        central = pytz.timezone('US/Central')
+        safe_date_time = datetime.datetime.now().strftime("%Y%m%d")
+        replaced_prompt = prompt.replace(" ", "_").replace("\n", "_")
+        safe_prompt = "".join(x for x in replaced_prompt if x.isalnum() or x == "_")[:90]
+        filename = f"{safe_date_time}_{safe_prompt}.png"
+        
+        # Save the image
+        if len(results.images) > 0:
+            image_path = os.path.join("", filename)  # Specify your directory
+            results.images[0].save(image_path)
+            print(f"#Image saved as {image_path}")
+            gr.File(image_path)
+            gr.Button(link=image_path)
+            # encoded_image = encode_image(image)
+            # html_link = f'<a href="data:image/png;base64,{encoded_image}" download="{filename}">Download Image</a>'
+            # gr.HTML(html_link)
+    except:
+        return results.images[0]
 
     return results.images[0] if len(results.images) > 0 else None
 
@@ -138,9 +230,10 @@ css = """
 }
 """
 with gr.Blocks(css=css) as demo:
+
     with gr.Column(elem_id="container"):
         gr.Markdown(
-            """## 🕹️ Stable Diffusion 1.5 - Real Time 🎨 Image Generation Using 🌐 Latent Consistency LoRAs""",
+            """4📝RT🖼️Images - 🕹️ Real Time 🎨 Image Generator Gallery 🌐""",
             elem_id="intro",
         )
         with gr.Row():
@@ -166,6 +259,12 @@ with gr.Blocks(css=css) as demo:
             label="Generated Images", show_label=False, elem_id="gallery"
         )
 
+        with gr.Row(variant="compact"):
+            # 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)
+
         # Advanced Generate Options
         with gr.Accordion("Advanced options", open=False):
             guidance = gr.Slider(
@@ -186,11 +285,9 @@ with gr.Blocks(css=css) as demo:
             
             ```py
             from diffusers import DiffusionPipeline, LCMScheduler
-
             pipe = DiffusionPipeline.from_pretrained("Lykon/dreamshaper-7").to("cuda") 
             pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
             pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5") #yes, it's a normal LoRA
-
             results = pipe(
                 prompt="ImageEditor",
                 num_inference_steps=4,
@@ -210,5 +307,11 @@ with gr.Blocks(css=css) as demo:
         steps.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
         seed.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
 
+        # Attach click event handlers to the buttons
+        save_all_button.click(save_all_button_click)
+        clear_all_button.click(clear_all_button_click)
+
+        
+
 demo.queue()
 demo.launch()