Update app.py

app.py

@@ -9,10 +9,10 @@ import torch
 from diffusers import DiffusionPipeline
 from PIL import Image
 
-# Create directories if they don't exist
-SAVE_DIR = "generated_images"
+# Create permanent storage directory
+SAVE_DIR = "/content/generated_images"
 if not os.path.exists(SAVE_DIR):
-    os.makedirs(SAVE_DIR)
+    os.makedirs(SAVE_DIR, exist_ok=True)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 repo_id = "black-forest-labs/FLUX.1-dev"
@@ -25,7 +25,7 @@ pipeline = pipeline.to(device)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def save_generated_image(image):
+def save_generated_image(image, prompt):
     # Generate unique filename with timestamp
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     unique_id = str(uuid.uuid4())[:8]
@@ -34,6 +34,12 @@ def save_generated_image(image):
     
     # Save the image
     image.save(filepath)
+    
+    # Save metadata
+    metadata_file = os.path.join(SAVE_DIR, "metadata.txt")
+    with open(metadata_file, "a", encoding="utf-8") as f:
+        f.write(f"{filename}|{prompt}|{timestamp}\n")
+    
     return filepath
 
 def load_generated_images():
@@ -47,6 +53,28 @@ 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",
+        "assets/cm2.webp",
+        "assets/cm3.webp",
+        "assets/cm4.webp",
+        "assets/cm5.webp",
+        "assets/cm6.webp",
+    ]
+    return predefined_images
+
 @spaces.GPU(duration=120)
 def inference(
     prompt: str,
@@ -73,23 +101,12 @@ def inference(
         joint_attention_kwargs={"scale": lora_scale},
     ).images[0]
     
-    # Save the generated image
-    save_generated_image(image)
+    # Save the generated image with the prompt
+    save_generated_image(image, prompt)
     
     # Return the image, seed, and updated gallery
     return image, seed, load_generated_images()
 
-def load_predefined_images():
-    predefined_images = [
-        "assets/cm1.webp",
-        "assets/cm2.webp",
-        "assets/cm3.webp",
-        "assets/cm4.webp",
-        "assets/cm5.webp",
-        "assets/cm6.webp",
-    ]
-    return predefined_images
-
 examples = [
     "Claude Monet's 1916 painting, Water Lilies, which is currently on display at the Metropolitan Museum of Art. The painting depicts a tranquil pond with water lilies floating on the surface, surrounded by lush green foliage and a variety of colorful flowers. The colors of the flowers range from bright pinks and purples to deep blues and greens, creating a peaceful and calming atmosphere. [trigger]",
     "Claude Monet's 1869 masterpiece, The Magpie, showcasing a snow-covered rural landscape at dawn. A single black magpie perches on a wooden gate, contrasting against the pristine white snow. The scene captures the subtle interplay of light and shadow on the snow's surface, with delicate blue-gray tones in the shadows and warm golden hints where sunlight touches the snow-laden branches. [trigger]",
@@ -106,6 +123,8 @@ footer {
 """
 
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
+    flagging_callback = ImageFlagging()
+
     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>')
 
@@ -185,7 +204,8 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
                 columns=6,
                 show_label=False,
                 value=load_generated_images(),
-            )
+                elem_id="generated_gallery"
+            ).style(grid=6)
 
     # Add sample gallery section at the bottom
     gr.Markdown("### Claude Monet Style Examples")
@@ -197,6 +217,7 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
         value=load_predefined_images()
     )
 
+    # Enable flagging for permanent storage
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=inference,
@@ -211,6 +232,12 @@ 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()