Update app.py

app.py

@@ -14,7 +14,8 @@ 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")
+# Change gallery path to user's home directory for persistence
+gallery_path = path.join(os.path.expanduser("~"), "ai_generated_images")
 os.environ["TRANSFORMERS_CACHE"] = cache_path
 os.environ["HF_HUB_CACHE"] = cache_path
 os.environ["HF_HOME"] = cache_path
@@ -110,19 +111,6 @@ footer {display: none !important}
     border-radius: 4px !important;
     transition: transform 0.2s;
 }
-/* Force gallery items to maintain aspect ratio */
-.gallery-item {
-    width: 100% !important;
-    aspect-ratio: 1 !important;
-    overflow: hidden !important;
-}
-.gallery-item img {
-    width: 100% !important;
-    height: 100% !important;
-    object-fit: cover !important;
-    border-radius: 4px;
-    transition: transform 0.2s;
-}
 .gallery-item img:hover {
     transform: scale(1.05);
 }
@@ -161,30 +149,60 @@ footer {display: none !important}
 
 def save_image(image):
     """Save the generated image and return the path"""
+    # Ensure gallery directory exists
+    os.makedirs(gallery_path, exist_ok=True)
+    
+    # Generate unique filename with timestamp and random suffix
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    filename = f"generated_{timestamp}.png"
+    random_suffix = os.urandom(4).hex()
+    filename = f"generated_{timestamp}_{random_suffix}.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
+    try:
+        if isinstance(image, Image.Image):
+            # Save with maximum quality
+            image.save(filepath, "PNG", quality=100)
+        else:
+            image = Image.fromarray(image)
+            image.save(filepath, "PNG", quality=100)
+        
+        # Verify the file was saved correctly
+        if not os.path.exists(filepath):
+            print(f"Warning: Failed to verify saved image at {filepath}")
+            return None
+            
+        return filepath
+    except Exception as e:
+        print(f"Error saving image: {str(e)}")
+        return None
 
 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]
+    try:
+        # Ensure gallery directory exists
+        os.makedirs(gallery_path, exist_ok=True)
+        
+        # Get all image files and sort by modification time
+        image_files = []
+        for f in os.listdir(gallery_path):
+            if f.lower().endswith(('.png', '.jpg', '.jpeg')):
+                full_path = os.path.join(gallery_path, f)
+                image_files.append((full_path, os.path.getmtime(full_path)))
+        
+        # Sort by modification time (newest first)
+        image_files.sort(key=lambda x: x[1], reverse=True)
+        
+        # Return only the file paths
+        return [f[0] for f in image_files]
+    except Exception as e:
+        print(f"Error loading gallery: {str(e)}")
+        return []
 
 # Create Gradio interface
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     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():
         with gr.Column(scale=3):
             prompt = gr.Textbox(
@@ -269,12 +287,14 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             """)
 
         with gr.Column(scale=4, elem_classes=["fixed-width"]):
-    # Current generated image
+            # Current generated image
             output = gr.Image(
                 label="Generated Image",
                 elem_id="output-image",
                 elem_classes=["output-image", "fixed-width"]
             )
+            
+            # Gallery of generated images
             gallery = gr.Gallery(
                 label="Generated Images Gallery",
                 show_label=True,
@@ -285,9 +305,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 object_fit="cover",
                 elem_classes=["gallery-container", "fixed-width"]
             )
-
-
-
             
             # Load existing gallery images on startup
             gallery.value = load_gallery()
@@ -296,21 +313,27 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     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"):
-            generated_image = pipe(
-                prompt=[prompt],
-                generator=torch.Generator().manual_seed(int(seed)),
-                num_inference_steps=int(steps),
-                guidance_scale=float(scales),
-                height=int(height),
-                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()
+            try:
+                generated_image = pipe(
+                    prompt=[prompt],
+                    generator=torch.Generator().manual_seed(int(seed)),
+                    num_inference_steps=int(steps),
+                    guidance_scale=float(scales),
+                    height=int(height),
+                    width=int(width),
+                    max_sequence_length=256
+                ).images[0]
+                
+                # Save the generated image
+                saved_path = save_image(generated_image)
+                if saved_path is None:
+                    print("Warning: Failed to save generated image")
+                
+                # Return both the generated image and updated gallery
+                return generated_image, load_gallery()
+            except Exception as e:
+                print(f"Error in image generation: {str(e)}")
+                return None, load_gallery()
     
     # Connect the generation button to both the image output and gallery update
     def update_seed():