Update app.py

app.py

@@ -4,45 +4,58 @@ import os
 import random
 from PIL import Image
 
-# Paths
-OUTPUT_DIR = "outputs"
-MODEL_PATH = "dress.pkl"  # Adjusted for local Hugging Face repo structure
+# === Setup Paths ===
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+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 the output directory exists
+# Ensure output directory exists
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-# Function to generate images using StyleGAN3
+# === Image Generation Function ===
 def generate_images():
-    command = f"python stylegan3/gen_images.py --outdir={OUTPUT_DIR} --trunc=1 --seeds='3-5,7,9,12-14,16-26,29,31,32,34,40,41' --network={MODEL_PATH}"
+    command = [
+        "python",
+        GEN_SCRIPT,
+        f"--outdir={OUTPUT_DIR}",
+        "--trunc=1",
+        "--seeds=3-5,7,9,12-14,16-26,29,31,32,34,40,41",
+        f"--network={MODEL_PATH}"
+    ]
+    
     try:
-        subprocess.run(command, shell=True, check=True)
+        result = subprocess.run(command, check=True, capture_output=True, text=True)
+        print("stdout:\n", result.stdout)
+        print("stderr:\n", result.stderr)
     except subprocess.CalledProcessError as e:
-        return f"Error generating images: {e}"
+        print("Image generation failed with error:")
+        print(e.stderr)
+        return f"Error generating images:\n{e.stderr}"
 
-# Function to select 5 random images
+# === Select Random Images from Output Folder ===
 def get_random_images():
     image_files = [f for f in os.listdir(OUTPUT_DIR) if f.endswith(".png")]
+
     if len(image_files) < 5:
         generate_images()
         image_files = [f for f in os.listdir(OUTPUT_DIR) if f.endswith(".png")]
-    
-    if len(image_files) == 0:
-        return ["No images found. Please try again."]
 
-    random_images = random.sample(image_files, min(5, len(image_files)))
-    return [os.path.join(OUTPUT_DIR, img) for img in random_images]
+    if not image_files:
+        return [Image.new("RGB", (256, 256), color="gray")] * 5  # fallback
 
+    random_images = random.sample(image_files, min(5, len(image_files)))
+    return [Image.open(os.path.join(OUTPUT_DIR, img)) for img in random_images]
 
-# Gradio function
+# === Gradio Interface ===
 def generate_and_display():
     generate_images()
     return get_random_images()
 
-# 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, rows=2)
+    output_gallery = gr.Gallery(label="Generated Designs", columns=5)
     generate_button.click(fn=generate_and_display, outputs=output_gallery)
 
 if __name__ == "__main__":