Update app.py

app.py

@@ -1,5 +1,7 @@
 import random
-
+import os
+import uuid
+from datetime import datetime
 import gradio as gr
 import numpy as np
 import spaces
@@ -7,6 +9,11 @@ import torch
 from diffusers import DiffusionPipeline
 from PIL import Image
 
+# Create directories if they don't exist
+SAVE_DIR = "generated_images"
+if not os.path.exists(SAVE_DIR):
+    os.makedirs(SAVE_DIR)
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 repo_id = "black-forest-labs/FLUX.1-dev"
 adapter_id = "openfree/claude-monet"
@@ -15,10 +22,30 @@ pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16
 pipeline.load_lora_weights(adapter_id)
 pipeline = pipeline.to(device)
 
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
+def save_generated_image(image):
+    # Generate unique filename with timestamp
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    unique_id = str(uuid.uuid4())[:8]
+    filename = f"{timestamp}_{unique_id}.png"
+    filepath = os.path.join(SAVE_DIR, filename)
+    
+    # Save the image
+    image.save(filepath)
+    return filepath
+
+def load_generated_images():
+    if not os.path.exists(SAVE_DIR):
+        return []
+    
+    # Load all images from the directory
+    image_files = [os.path.join(SAVE_DIR, f) for f in os.listdir(SAVE_DIR) 
+                  if f.endswith(('.png', '.jpg', '.jpeg', '.webp'))]
+    # Sort by creation time (newest first)
+    image_files.sort(key=lambda x: os.path.getctime(x), reverse=True)
+    return image_files
 
 @spaces.GPU(duration=120)
 def inference(
@@ -45,8 +72,12 @@ def inference(
         generator=generator,
         joint_attention_kwargs={"scale": lora_scale},
     ).images[0]
-
-    return image, seed
+    
+    # Save the generated image
+    save_generated_image(image)
+    
+    # Return the image, seed, and updated gallery
+    return image, seed, load_generated_images()
 
 def load_predefined_images():
     predefined_images = [
@@ -74,84 +105,97 @@ footer {
 }
 """
 
-with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css
-) as demo:
+with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) 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>')
 
-    
-    with gr.Column(elem_id="col-container"):
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-
-            run_button = gr.Button("Run", scale=0)
-
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=42,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=768,
-                )
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=3.5,
+    with gr.Tabs() as tabs:
+        with gr.Tab("Generation"):
+            with gr.Column(elem_id="col-container"):
+                with gr.Row():
+                    prompt = gr.Text(
+                        label="Prompt",
+                        show_label=False,
+                        max_lines=1,
+                        placeholder="Enter your prompt",
+                        container=False,
+                    )
+                    run_button = gr.Button("Run", scale=0)
+
+                result = gr.Image(label="Result", show_label=False)
+
+                with gr.Accordion("Advanced Settings", open=False):
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        step=1,
+                        value=42,
+                    )
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+                    with gr.Row():
+                        width = gr.Slider(
+                            label="Width",
+                            minimum=256,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=32,
+                            value=1024,
+                        )
+                        height = gr.Slider(
+                            label="Height",
+                            minimum=256,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=32,
+                            value=768,
+                        )
+
+                    with gr.Row():
+                        guidance_scale = gr.Slider(
+                            label="Guidance scale",
+                            minimum=0.0,
+                            maximum=10.0,
+                            step=0.1,
+                            value=3.5,
+                        )
+                        num_inference_steps = gr.Slider(
+                            label="Number of inference steps",
+                            minimum=1,
+                            maximum=50,
+                            step=1,
+                            value=30,
+                        )
+                        lora_scale = gr.Slider(
+                            label="LoRA scale",
+                            minimum=0.0,
+                            maximum=1.0,
+                            step=0.1,
+                            value=1.0,
+                        )
+
+                gr.Examples(
+                    examples=examples,
+                    inputs=[prompt],
+                    outputs=[result, seed],
                 )
 
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=30,
-                )
-
-                lora_scale = gr.Slider(
-                    label="LoRA scale",
-                    minimum=0.0,
-                    maximum=1.0,
-                    step=0.1,
-                    value=1.0,
-                )
+        with gr.Tab("Gallery"):
+            generated_gallery = gr.Gallery(
+                label="Generated Images",
+                columns=6,
+                show_label=False,
+                value=load_generated_images(),
+            )
 
-        gr.Examples(
-            examples=examples,
-            inputs=[prompt],
-            outputs=[result, seed],
-        )
+    # Add sample gallery section at the bottom
+    gr.Markdown("### Claude Monet Style Examples")
+    predefined_gallery = gr.Gallery(
+        label="Sample Images", 
+        columns=3, 
+        rows=2,
+        show_label=False, 
+        value=load_predefined_images()
+    )
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
@@ -166,17 +210,7 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css
             num_inference_steps,
             lora_scale,
         ],
-        outputs=[result, seed],
-    )
-    
-    # Add gallery section at the bottom
-    gr.Markdown("### Claude Monet Style Examples")
-    predefined_gallery = gr.Gallery(
-        label="Sample Images", 
-        columns=3, 
-        rows=2,
-        show_label=False, 
-        value=load_predefined_images()
+        outputs=[result, seed, generated_gallery],
     )
 
 demo.queue()