Update app.py

app.py

@@ -3,19 +3,27 @@ import argparse
 import os
 import time
 from os import path
+import shutil
+from datetime import datetime
 from safetensors.torch import load_file
 from huggingface_hub import hf_hub_download
 import gradio as gr
 import torch
 from diffusers import FluxPipeline
+from PIL import Image
 
 # Setup and initialization code
 cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
+gallery_path = path.join(path.dirname(path.abspath(__file__)), "gallery")
 os.environ["TRANSFORMERS_CACHE"] = cache_path
 os.environ["HF_HUB_CACHE"] = cache_path
 os.environ["HF_HOME"] = cache_path
 torch.backends.cuda.matmul.allow_tf32 = True
 
+# Create gallery directory if it doesn't exist
+if not path.exists(gallery_path):
+    os.makedirs(gallery_path, exist_ok=True)
+
 class timer:
     def __init__(self, method_name="timed process"):
         self.method = method_name
@@ -58,11 +66,50 @@ footer {display: none !important}
     -webkit-background-clip: text;
     -webkit-text-fill-color: transparent;
 }
+.gallery-container {
+    display: grid;
+    grid-template-columns: repeat(auto-fill, minmax(150px, 1fr));
+    gap: 10px;
+    padding: 10px;
+    background: rgba(255, 255, 255, 0.05);
+    border-radius: 8px;
+    margin-top: 10px;
+}
+.gallery-image {
+    width: 100%;
+    aspect-ratio: 1;
+    object-fit: cover;
+    border-radius: 4px;
+    transition: transform 0.2s;
+}
+.gallery-image:hover {
+    transform: scale(1.05);
+}
 """
 
+def save_image(image):
+    """Save the generated image and return the path"""
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"generated_{timestamp}.png"
+    filepath = os.path.join(gallery_path, filename)
+    
+    if isinstance(image, Image.Image):
+        image.save(filepath)
+    else:
+        image = Image.fromarray(image)
+        image.save(filepath)
+    
+    return filepath
+
+def load_gallery():
+    """Load all images from the gallery directory"""
+    image_files = [f for f in os.listdir(gallery_path) if f.endswith(('.png', '.jpg', '.jpeg'))]
+    image_files.sort(reverse=True)  # Most recent first
+    return [os.path.join(gallery_path, f) for f in image_files]
+
 # Create Gradio interface
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
-    gr.HTML('<div class="title">Flux 8Step LoRA: Image Generator</div>')
+    gr.HTML('<div class="title">AI Image Generator</div>')
     gr.HTML('<div style="text-align: center; margin-bottom: 2em; color: #666;">Create stunning images from your descriptions</div>')
     
     with gr.Row():
@@ -140,24 +187,32 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                         <p style="font-weight: bold; margin: 0 0 0.5em 0;">🌿 Macro Nature</p>
                         <p style="margin: 0; font-style: italic;">"Extreme macro photography of a dewdrop on a butterfly wing, rainbow light refraction, crystalline clarity, intricate wing scales visible, natural bokeh background, professional studio lighting, shot with Canon MP-E 65mm lens"</p>
                     </div>
-                    <h4 style="margin: 1em 0 0.5em 0;">Tips for best results:</h4>
-                    <ul style="margin: 0; padding-left: 1.2em;">
-                        <li>Be specific in your descriptions</li>
-                        <li>Include details about style, lighting, and mood</li>
-                        <li>Reference specific artists or techniques</li>
-                        <li>Experiment with different guidance scales</li>
-                    </ul>
                 </div>
             """)
         
         with gr.Column(scale=4):
+            # Current generated image
             output = gr.Image(label="Generated Image")
+            
+            # Gallery of generated images
+            gallery = gr.Gallery(
+                label="Generated Images Gallery",
+                show_label=True,
+                elem_id="gallery",
+                columns=[4],
+                rows=[2],
+                height="auto",
+                object_fit="contain"
+            )
+            
+            # Load existing gallery images on startup
+            gallery.value = load_gallery()
     
     @spaces.GPU
-    def process_image(height, width, steps, scales, prompt, seed):
+    def process_and_save_image(height, width, steps, scales, prompt, seed):
         global pipe
         with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
-            return pipe(
+            generated_image = pipe(
                 prompt=[prompt],
                 generator=torch.Generator().manual_seed(int(seed)),
                 num_inference_steps=int(steps),
@@ -166,11 +221,18 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 width=int(width),
                 max_sequence_length=256
             ).images[0]
+            
+            # Save the generated image
+            save_image(generated_image)
+            
+            # Return both the generated image and updated gallery
+            return generated_image, load_gallery()
     
+    # Connect the generation button to both the image output and gallery update
     generate_btn.click(
-        process_image,
+        process_and_save_image,
         inputs=[height, width, steps, scales, prompt, seed],
-        outputs=output
+        outputs=[output, gallery]
     )
 
 if __name__ == "__main__":