Update app.py

app.py

@@ -57,11 +57,14 @@
 #     demo.launch()
 
 
+
+
 import gradio as gr
 import subprocess
 import os
 import random
 from PIL import Image
+import shutil
 
 # === Setup Paths ===
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
@@ -99,58 +102,33 @@ def get_random_images():
         generate_images()
         image_files = [f for f in os.listdir(OUTPUT_DIR) if f.endswith(".png")]
     random_images = random.sample(image_files, min(10, len(image_files)))
-    images = [Image.open(os.path.join(OUTPUT_DIR, img)) for img in random_images]
-    return images, random_images
+    return [Image.open(os.path.join(OUTPUT_DIR, img)) for img in random_images], image_files
 
 # === Save Image Function ===
-def save_image(image_name):
-    image_path = os.path.join(OUTPUT_DIR, image_name)
-    return image_path
+def save_image(image_filename):
+    image_path = os.path.join(OUTPUT_DIR, image_filename)
+    save_path = os.path.join(BASE_DIR, "saved_images", image_filename)
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    shutil.copy(image_path, save_path)
+    return f"Image saved as {save_path}"
 
 # === Gradio Interface ===
 def generate_and_display():
-    images, image_names = get_random_images()
-    return images, image_names, None  # Clear selected image path on generate
+    images, filenames = get_random_images()
+    buttons = [gr.Button("Save", elem_id=filename) for filename in filenames]
+    return images, buttons
+
+def save_button_clicked(image_filename):
+    return save_image(image_filename)
 
 with gr.Blocks() as demo:
     gr.Markdown("# 🎨 AI-Generated Clothing Designs - Dresses")
     generate_button = gr.Button("Generate New Designs")
-    output_gallery = gr.Gallery(label="Generated Designs", columns=5, height=300)
-    # Variable to hold image filenames aligned with gallery images
-    image_names = gr.State()
-    
-    saved_image_path = gr.Textbox(label="Save Image Path", visible=False)
-    
-    generate_button.click(fn=generate_and_display, inputs=None, outputs=[output_gallery, image_names, saved_image_path])
-    
-    # When user selects image from gallery, output selected image filename and show the save button
-    def on_select(idx, names):
-        if names and idx is not None and 0 <= idx < len(names):
-            return names[idx]
-        return None
-    
-    selected_image = gr.Textbox(label="Selected Image", visible=False)
-    output_gallery.select(fn=on_select, inputs=[gr.Gallery.index, image_names], outputs=selected_image)
-    
-    # Save Button, outputs a downloadable href link to the image path
-    with gr.Row(visible=False) as save_row:
-        save_button = gr.Button("Save Selected Image")
-        download_link = gr.File(label="Download Saved Image")
-    
-    def show_save_button(selected_img_path):
-        # Show save button only if an image is selected
-        return gr.Row.update(visible=bool(selected_img_path))
-    
-    selected_image.change(fn=show_save_button, inputs=selected_image, outputs=save_row)
-    
-    def save_selected_image(selected_img_path):
-        if selected_img_path and os.path.exists(selected_img_path):
-            # Return path to file for download
-            return selected_img_path
-        return None
-    
-    save_button.click(fn=save_selected_image, inputs=selected_image, outputs=download_link)
+    output_gallery = gr.Gallery(label="Generated Designs", columns=5)
+    generate_button.click(fn=generate_and_display, outputs=output_gallery)
+
+    # Add save buttons to each image dynamically
+    output_gallery.select(fn=save_button_clicked, inputs=output_gallery, outputs=None)
 
 if __name__ == "__main__":
     demo.launch()
-