Update app.py

app.py

@@ -56,7 +56,6 @@
 # if __name__ == "__main__":
 #     demo.launch()
 
-
 import gradio as gr
 import subprocess
 import os
@@ -70,7 +69,6 @@ GEN_SCRIPT = os.path.join(BASE_DIR, "stylegan3", "gen_images.py")
 OUTPUT_DIR = os.path.join(BASE_DIR, "outputs")
 MODEL_PATH = os.path.join(BASE_DIR, "dress.pkl")
 
-# Ensure output directory exists
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 # === Image Generation Function ===
@@ -83,52 +81,52 @@ def generate_images():
         "--seeds=3-5,7,9,12-14,16-26,29,31,32,34,40,41",
         f"--network={MODEL_PATH}"
     ]
-    
     try:
-        result = subprocess.run(command, check=True, capture_output=True, text=True)
-        print("stdout:\n", result.stdout)
-        print("stderr:\n", result.stderr)
+        subprocess.run(command, check=True, capture_output=True, text=True)
     except subprocess.CalledProcessError as e:
-        print("Image generation failed with error:")
-        print(e.stderr)
         return f"Error generating images:\n{e.stderr}"
 
-# === Select Random Images from Output Folder ===
+# === Get 10 Random Images ===
 def get_random_images():
     image_files = [f for f in os.listdir(OUTPUT_DIR) if f.endswith(".png")]
     if len(image_files) < 10:
         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
+    selected_files = random.sample(image_files, min(10, len(image_files)))
+    return selected_files
 
 # === Save Image Function ===
-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, filenames = get_random_images()
-    save_buttons = [gr.Button("Save", elem_id=filename) for filename in filenames]
-    return images, save_buttons, filenames
-
-def save_button_clicked(image_filename):
-    return save_image(image_filename)
-
+def save_image(filename):
+    src_path = os.path.join(OUTPUT_DIR, filename)
+    dest_dir = os.path.join(BASE_DIR, "saved_images")
+    os.makedirs(dest_dir, exist_ok=True)
+    dest_path = os.path.join(dest_dir, filename)
+    shutil.copy(src_path, dest_path)
+    return f"✅ Saved as {dest_path}"
+
+# === UI ===
 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)
-    
-    generate_button.click(fn=generate_and_display, outputs=[output_gallery, generate_button])
+    gr.Markdown("# 👗 AI-Generated Clothing Designs (Dresses)")
+
+    generate_btn = gr.Button("✨ Generate New Designs")
+    image_rows = gr.Column()
 
-    # Create dynamic save buttons for each generated image
-    with gr.Row() as button_row:
-        for filename in []:
-            gr.Button("Save", elem_id=filename).click(fn=save_button_clicked, inputs=filename, outputs=None)
-        
+    def generate_ui():
+        files = get_random_images()
+        elements = []
+        for f in files:
+            img_path = os.path.join(OUTPUT_DIR, f)
+            with gr.Row():
+                image = gr.Image(img_path, label=f, show_label=False)
+                save_btn = gr.Button("💾 Save This Image")
+
+                # Use a closure to capture the filename
+                save_btn.click(fn=save_image, inputs=gr.Textbox(value=f, visible=False), outputs=gr.Textbox())
+
+        return
+
+    generate_btn.click(fn=generate_ui, outputs=[])
+
+if __name__ == "__main__":
     demo.launch()
+