Update app.py

app.py

@@ -9,7 +9,7 @@ import re
 import random
 import torch
 import time
-import shutil
+import shutil  # Added for zip functionality
 import zipfile
 from PIL import Image
 from io import BytesIO
@@ -23,6 +23,7 @@ except:
 SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", None)
 TORCH_COMPILE = os.environ.get("TORCH_COMPILE", None)
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
+# check if MPS is available OSX only M1/M2/M3 chips
 mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
 xpu_available = hasattr(torch, "xpu") and torch.xpu.is_available()
 device = torch.device(
@@ -31,73 +32,79 @@ device = torch.device(
 torch_device = device
 torch_dtype = torch.float16
 
-# CSS definition
-css = """
-#container{
-    margin: 0 auto;
-    max-width: 40rem;
-}
-#intro{
-    max-width: 100%;
-    text-align: center;
-    margin: 0 auto;
-}
-"""
-
+# 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()
-    img_files = [file for file in os.listdir(base_dir) if file.lower().endswith((".png", ".jpg", ".jpeg"))]
+    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)
-
-def save_all_images(images):
-    if len(images) == 0:
-        return None, None
-    
+  
+# 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_and_history_{timestamp}.zip"
-    
+    zip_filename = f"images_{timestamp}.zip"
+    print('Creating file ' + zip_filename)
     with zipfile.ZipFile(zip_filename, 'w') as zipf:
-        # Add image files
-        for file in images:
+        for file in files:
             zipf.write(file, os.path.basename(file))
-        
-        # Add prompt history file
-        if os.path.exists("prompt_history.txt"):
-            zipf.write("prompt_history.txt")
-    
-    # Generate download link
-    zip_base64 = encode_file_to_base64(zip_filename)
-    download_link = f'<a href="data:application/zip;base64,{zip_base64}" download="{zip_filename}">Download All (Images & History)</a>'
+            print('added:' + file)
+    return zip_filename
+
+def get_zip_download_link(zip_file):
+    """
+    Generate a link to download the zip file.
+    """
+    zip_base64 = encode_file_to_base64(zip_file)  # Encode the zip file to base64
+    href = f'<a href="data:application/zip;base64,{zip_base64}" download="{zip_file}">Download All</a>'
+    return href
     
+def save_all_images(images):
+    if len(images) == 0:
+        return None, None
+    zip_filename = create_zip_of_files(images)  # Create a zip file from the list of image files
+    print(f"Zip file created: {zip_filename}")
+    download_link = get_zip_download_link(zip_filename)
     return zip_filename, download_link
-
+        
 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:
-        return gr.HTML(download_link)
+        gr.HTML(download_link)
 
+
+# Function to handle "Clear All" button click
 def clear_all_button_click():
     clear_all_images()
 
@@ -120,6 +127,7 @@ pipe.to(device=torch_device, dtype=torch_dtype).to(device)
 pipe.unet.to(memory_format=torch.channels_last)
 pipe.set_progress_bar_config(disable=True)
 
+# check if computer has less than 64GB of RAM using sys or os
 if psutil.virtual_memory().total < 64 * 1024**3:
     pipe.enable_attention_slicing()
 
@@ -128,29 +136,28 @@ if TORCH_COMPILE:
     pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
     pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0)
 
+# Load LCM LoRA
 pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
 pipe.fuse_lora()
 
 def safe_filename(text):
+    """Generate a safe filename from a string."""
     safe_text = re.sub(r'\W+', '_', text)
     timestamp = datetime.datetime.now().strftime("%Y%m%d")
     return f"{safe_text}_{timestamp}.png"
-
+    
 def encode_image(image):
+    """Encode image to base64."""
     buffered = BytesIO()
+    #image.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode()
 
 def fake_gan():
-    base_dir = os.getcwd()
-    img_files = [file for file in os.listdir(base_dir) if file.lower().endswith((".png", ".jpg", ".jpeg"))]
+    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"
     images = [(random.choice(img_files), os.path.splitext(file)[0]) for file in img_files]
     return images
-
-def save_prompt_to_history(prompt):
-    with open("prompt_history.txt", "a") as f:
-        timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        f.write(f"{timestamp}: {prompt}\n")
-
+    
 def predict(prompt, guidance, steps, seed=1231231):
     generator = torch.manual_seed(seed)
     last_time = time.time()
@@ -161,13 +168,10 @@ def predict(prompt, guidance, steps, seed=1231231):
         guidance_scale=guidance,
         width=512,
         height=512,
+        # original_inference_steps=params.lcm_steps,
         output_type="pil",
     )
     print(f"Pipe took {time.time() - last_time} seconds")
-    
-    # Save prompt to history
-    save_prompt_to_history(prompt)
-    
     nsfw_content_detected = (
         results.nsfw_content_detected[0]
         if "nsfw_content_detected" in results
@@ -183,24 +187,35 @@ def predict(prompt, guidance, steps, seed=1231231):
         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)
+            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
 
-def read_prompt_history():
-    if os.path.exists("prompt_history.txt"):
-        with open("prompt_history.txt", "r") as f:
-            return f.read()
-    return "No prompts yet."
 
+css = """
+#container{
+    margin: 0 auto;
+    max-width: 40rem;
+}
+#intro{
+    max-width: 100%;
+    text-align: center;
+    margin: 0 auto;
+}
+"""
 with gr.Blocks(css=css) as demo:
+
     with gr.Column(elem_id="container"):
         gr.Markdown(
             """4📝RT🖼️Images - 🕹️ Real Time 🎨 Image Generator Gallery 🌐""",
@@ -215,8 +230,10 @@ with gr.Blocks(css=css) as demo:
 
                 gr.Button("Download", link="/file=all_files.zip")
     
+        # Image Result from last prompt
         image = gr.Image(type="filepath")
 
+        # Gallery of Generated Images with Image Names in Random Set to Download
         with gr.Row(variant="compact"):
             text = gr.Textbox(
                 label="Image Sets",
@@ -231,9 +248,12 @@ with gr.Blocks(css=css) as demo:
         )
 
         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(
                 label="Guidance", minimum=0.0, maximum=5, value=0.3, step=0.001
@@ -243,54 +263,46 @@ with gr.Blocks(css=css) as demo:
                 randomize=True, minimum=0, maximum=12013012031030, label="Seed", step=1
             )
 
-        with gr.Accordion("Prompt History", open=False):
-            prompt_history = gr.Textbox(label="Prompt History", lines=10, max_lines=20, interactive=False)
-
+        # Diffusers
         with gr.Accordion("Run with diffusers"):
             gr.Markdown(
                 """## Running LCM-LoRAs it with `diffusers`
-                ```bash
-                pip install diffusers==0.23.0
-                ```
-                
-                ```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,
-                    guidance_scale=0.0,
-                )
-                results.images[0]
-                ```
-                """
+            ```bash
+            pip install diffusers==0.23.0
+            ```
+            
+            ```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,
+                guidance_scale=0.0,
+            )
+            results.images[0]
+            ```
+            """
             )
 
-        with gr.Column():
-            file_obj = gr.File(label="Input File")
-            input = file_obj
-
-    inputs = [prompt, guidance, steps, seed]
-    generate_bt.click(fn=predict, inputs=inputs, outputs=[image, prompt_history], show_progress=False)
-    btn.click(fake_gan, None, gallery)
-    prompt.submit(fn=predict, inputs=inputs, outputs=[image, prompt_history], show_progress=False)
-    guidance.change(fn=predict, inputs=inputs, outputs=[image, prompt_history], show_progress=False)
-    steps.change(fn=predict, inputs=inputs, outputs=[image, prompt_history], show_progress=False)
-    seed.change(fn=predict, inputs=inputs, outputs=[image, prompt_history], show_progress=False)
-
-    def update_prompt_history():
-        return read_prompt_history()
+        # Function IO Eventing and Controls
+        inputs = [prompt, guidance, steps, seed]
+        generate_bt.click(fn=predict, inputs=inputs, outputs=image, show_progress=False)
+        btn.click(fake_gan, None, gallery)
+        prompt.input(fn=predict, inputs=inputs, outputs=image, show_progress=False)
+        guidance.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
+        steps.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
+        seed.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
 
-    generate_bt.click(fn=update_prompt_history, outputs=prompt_history)
-    prompt.submit(fn=update_prompt_history, outputs=prompt_history)
+        # Attach click event handlers to the buttons
+        save_all_button.click(save_all_button_click)
 
-    save_all_button.click(
-        fn=lambda: save_all_images([f for f in os.listdir() if f.lower().endswith((".png", ".jpg", ".jpeg"))]),
-        outputs=[gr.File(), gr.HTML()]
-    )
-    clear_all_button.click(clear_all_button_click)
+        with gr.Column():
+            file_obj = gr.File(label="Input File")
+            input= file_obj
+                    
+        clear_all_button.click(clear_all_button_click)
 
 demo.queue()
-demo.launch(allowed_paths=["/"])
+demo.launch(allowed_paths=["/"])