Update app.py

app.py

@@ -3,10 +3,12 @@ import requests
 from PIL import Image
 from io import BytesIO
 from tqdm import tqdm
+import os
 import time
 
 repo = "artificialguybr/TshirtDesignRedmond-V2"
 
+# Generate design based on prompts
 def infer(color_prompt, phone_type_prompt, design_prompt):
     prompt = (
         f"A single vertical {color_prompt} colored {phone_type_prompt} back cover featuring a bold {design_prompt} design on the front, hanging on the plain wall. The soft light and shadows, creating a striking contrast against the minimal background, evoking modern sophistication."
@@ -41,12 +43,26 @@ def infer(color_prompt, phone_type_prompt, design_prompt):
         else:
             raise Exception(f"API Error: {response.status_code}")
 
-def save_design(image_array):
-    image = Image.fromarray(image_array)
-    file_path = "saved_design.png"
+# Save the design
+def save_design(image):
+    file_path = "saved_designs/saved_design.png"
+    os.makedirs(os.path.dirname(file_path), exist_ok=True)
     image.save(file_path)
     return f"Design saved as {file_path}!"
 
+# Fetch saved designs for display on the Home page
+def fetch_saved_designs():
+    saved_designs_dir = "saved_designs"
+    if not os.path.exists(saved_designs_dir):
+        return ["No saved designs found."]
+    
+    designs = []
+    for file_name in os.listdir(saved_designs_dir):
+        if file_name.endswith((".png", ".jpg", ".jpeg")):
+            designs.append(os.path.join(saved_designs_dir, file_name))
+    return designs
+
+# Custom CSS for animations
 custom_css = """
 body {
     font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
@@ -73,19 +89,10 @@ body::before {
 }
 """
 
+# JavaScript for text-to-speech and particles
 custom_js = """
 <script>
 document.addEventListener('DOMContentLoaded', function () {
-    // Add particles
-    const particlesContainer = document.createElement('div');
-    particlesContainer.classList.add('particles');
-    for (let i = 0; i < 5; i++) {
-        const particle = document.createElement('div');
-        particle.classList.add('particle');
-        particlesContainer.appendChild(particle);
-    }
-    document.body.appendChild(particlesContainer);
-
     // Text-to-speech functionality
     function speak(text) {
         const synth = window.speechSynthesis;
@@ -102,6 +109,7 @@ document.addEventListener('DOMContentLoaded', function () {
 </script>
 """
 
+# Gradio interface
 with gr.Blocks(css=custom_css) as interface:
     gr.HTML(custom_js)
     gr.Markdown("# **AI Phone Cover Designer**")
@@ -110,8 +118,10 @@ with gr.Blocks(css=custom_css) as interface:
     # Navigation Tabs
     with gr.Tabs():
         with gr.Tab("Home"):
-            gr.Markdown("Welcome to the **AI Phone Cover Designer**! Use the 'Design' tab to start creating custom designs.")
-        
+            gr.Markdown("Welcome to the **AI Phone Cover Designer**! Below are some of your saved designs.")
+            saved_images_gallery = gr.Gallery(label="Saved Designs")
+            saved_images_gallery.update(value=fetch_saved_designs())
+
         with gr.Tab("Design"):
             with gr.Row():
                 with gr.Column(scale=1):
@@ -138,7 +148,7 @@ with gr.Blocks(css=custom_css) as interface:
             )
 
         with gr.Tab("About"):
-            gr.Markdown("""
+            gr.Markdown("""            
             ## About AI Phone Cover Maker
             The **AI Phone Cover Maker** is a cutting-edge tool designed to help users create personalized phone cover designs quickly and easily. 
             Powered by AI, it uses advanced image generation techniques to craft unique, high-quality designs for any mobile device.
@@ -151,4 +161,5 @@ with gr.Blocks(css=custom_css) as interface:
             Start designing today and bring your creative ideas to life!
             """)
 
+# Launch the app
 interface.launch(debug=True)