Update app.py

app.py

@@ -4,6 +4,8 @@ import random
 import spaces
 import os
 import torch
+import re
+from PIL import Image
 
 from diffusers import DiffusionPipeline, AutoencoderTiny
 from huggingface_hub import login
@@ -30,19 +32,25 @@ pipe = DiffusionPipeline.from_pretrained(
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 pipe.load_lora_weights("ZennyKenny/flux_lora_natalie-diffusion")
 
+# Ensure image_preview dir exists
+os.makedirs("image_preview", exist_ok=True)
+
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
+def sanitize_filename(name):
+    return re.sub(r"[^a-zA-Z0-9_-]", "_", name)[:80]
+
 @spaces.GPU(duration=75)
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
 
-    prompt = f"XTON {prompt}"
+    full_prompt = f"XTON {prompt}"
 
     for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
-        prompt=prompt,
+        prompt=full_prompt,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
         width=width,
@@ -50,7 +58,14 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidan
         generator=generator,
         output_type="pil",
     ):
-        yield img, seed
+        # Save image to /image_preview
+        safe_name = sanitize_filename(prompt)
+        img_path = f"image_preview/{safe_name}_{seed}.jpg"
+        img.convert("RGB").save(img_path, "JPEG", quality=60)
+
+        # Return image and updated list of previews
+        previews = [f"image_preview/{f}" for f in sorted(os.listdir("image_preview")) if f.endswith(".jpg")]
+        return img, seed, previews
 
 examples = [
     "a man walking in the forest",
@@ -62,14 +77,12 @@ css = """
 #left-column {
     padding: 1rem;
 }
-
 #right-column {
     display: flex;
     align-items: center;
     justify-content: center;
     padding: 1rem;
 }
-
 #generated-image > img {
     border-radius: 12px;
     box-shadow: 0px 0px 12px rgba(0, 0, 0, 0.2);
@@ -111,14 +124,13 @@ Generate images in the surreal style of artist [Natalie Kav](https://www.behance
                     guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=15, step=0.1, value=3.5)
                     num_inference_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=50, step=1, value=28)
 
-            # Use hidden result for examples only
             result_example = gr.Image(visible=False)
 
             gr.Examples(
                 examples=examples,
                 fn=infer,
                 inputs=[prompt],
-                outputs=[result_example, seed],
+                outputs=[result_example, seed, gr.Gallery(visible=False)],
                 cache_examples=True,
                 cache_mode="lazy"
             )
@@ -126,11 +138,16 @@ Generate images in the surreal style of artist [Natalie Kav](https://www.behance
         with gr.Column(scale=1, elem_id="right-column"):
             result = gr.Image(label="", show_label=False, elem_id="generated-image")
 
+    gr_state = gr.State([])  # internal list of previews
+    with gr.Column():
+        gr.Markdown("<h3 style='text-align:center;'>Generated Images Preview</h3>")
+        gallery = gr.Gallery(label="", columns=4, height="auto", object_fit="cover")
+
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
         inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs=[result, seed],
+        outputs=[result, seed, gallery],
     )
 
 if __name__ == "__main__":