Update app.py

app.py

@@ -9,8 +9,8 @@ import torch
 from diffusers import DiffusionPipeline
 from PIL import Image
 
-# Create permanent storage directory in the current working directory
-SAVE_DIR = os.path.join(os.getcwd(), "generated_images")
+# Create permanent storage directory
+SAVE_DIR = "saved_images"  # Gradio will handle the persistence
 if not os.path.exists(SAVE_DIR):
     os.makedirs(SAVE_DIR, exist_ok=True)
 
@@ -53,17 +53,6 @@ def load_generated_images():
     image_files.sort(key=lambda x: os.path.getctime(x), reverse=True)
     return image_files
 
-class ImageFlagging(gr.FlaggingCallback):
-    def setup(self, components, flagging_dir: str):
-        self.components = components
-        self.flagging_dir = SAVE_DIR
-        
-    def flag(self, flag_data, flag_option=None, flag_index=None, username=None) -> int:
-        """Save image and metadata permanently"""
-        image, prompt = flag_data
-        filepath = save_generated_image(image, prompt)
-        return 0
-
 def load_predefined_images():
     predefined_images = [
         "assets/cm1.webp",
@@ -101,8 +90,8 @@ def inference(
         joint_attention_kwargs={"scale": lora_scale},
     ).images[0]
     
-    # Save the generated image with the prompt
-    save_generated_image(image, prompt)
+    # Save the generated image
+    filepath = save_generated_image(image, prompt)
     
     # Return the image, seed, and updated gallery
     return image, seed, load_generated_images()
@@ -122,9 +111,7 @@ footer {
 }
 """
 
-with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
-    flagging_callback = ImageFlagging()
-
+with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css, analytics_enabled=False) as demo:
     gr.HTML('<div class="title"> Claude Monet STUDIO </div>')
     gr.HTML('<div class="title">😄Image to Video Explore: <a href="https://huggingface.co/spaces/ginigen/theater" target="_blank">https://huggingface.co/spaces/ginigen/theater</a></div>')
 
@@ -199,25 +186,37 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
                 )
 
         with gr.Tab("Gallery"):
+            gallery_header = gr.Markdown("### Generated Images Gallery")
             generated_gallery = gr.Gallery(
                 label="Generated Images",
                 columns=6,
                 show_label=False,
                 value=load_generated_images(),
-                elem_id="generated_gallery"
+                elem_id="generated_gallery",
+                height="auto"
             )
+            refresh_btn = gr.Button("🔄 Refresh Gallery")
 
     # Add sample gallery section at the bottom
     gr.Markdown("### Claude Monet Style Examples")
     predefined_gallery = gr.Gallery(
         label="Sample Images", 
-        columns=3, 
+        columns=3,
         rows=2,
         show_label=False, 
         value=load_predefined_images()
     )
 
-    # Enable flagging for permanent storage
+    # Event handlers
+    def refresh_gallery():
+        return load_generated_images()
+
+    refresh_btn.click(
+        fn=refresh_gallery,
+        inputs=None,
+        outputs=generated_gallery,
+    )
+
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=inference,
@@ -232,12 +231,6 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
             lora_scale,
         ],
         outputs=[result, seed, generated_gallery],
-    ).then(
-        fn=None,
-        inputs=[result, prompt],
-        outputs=None,
-        _js="(res, prompt) => clearInterval(window.GalleryIntervalID)",
-        flagging_callback=flagging_callback
     )
 
 demo.queue()