Update app.py

app.py

@@ -56,6 +56,7 @@
 # if __name__ == "__main__":
 #     demo.launch()
 
+
 import gradio as gr
 import subprocess
 import os
@@ -98,7 +99,8 @@ 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)))
-    return [Image.open(os.path.join(OUTPUT_DIR, img)) for img in random_images], random_images
+    images = [Image.open(os.path.join(OUTPUT_DIR, img)) for img in random_images]
+    return images, random_images
 
 # === Save Image Function ===
 def save_image(image_name):
@@ -108,27 +110,47 @@ def save_image(image_name):
 # === Gradio Interface ===
 def generate_and_display():
     images, image_names = get_random_images()
-    return images, image_names
+    return images, image_names, None  # Clear selected image path on generate
 
 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).style(height=300)
+    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
     
-    # Output for image names
-    image_names = gr.Variable()
+    selected_image = gr.Textbox(label="Selected Image", visible=False)
+    output_gallery.select(fn=on_select, inputs=[gr.Gallery.index, image_names], outputs=selected_image)
     
-    generate_button.click(fn=generate_and_display, outputs=[output_gallery, image_names])
+    # 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")
     
-    # Add a button to save the selected image
-    save_button = gr.Button("Save Selected 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))
     
-    def save_selected_image(selected_image_name):
-        if selected_image_name:
-            return save_image(selected_image_name)
+    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
-
-    output_gallery.select(fn=save_selected_image, inputs=image_names, outputs=save_button)
+    
+    save_button.click(fn=save_selected_image, inputs=selected_image, outputs=download_link)
 
 if __name__ == "__main__":
     demo.launch()
+