add gallery for iterations

app.py

@@ -2,6 +2,24 @@ import gradio as gr
 from main import main
 from arguments import parse_args
 import os
+import glob
+
+def list_iter_images(save_dir):
+    # Specify the image extensions you want to search for
+    image_extensions = ['jpg', 'jpeg', 'png', 'gif', 'bmp']  # Add more if needed
+
+    # Create a list to store the image file paths
+    image_paths = []
+
+    # Iterate through the specified image extensions and get the file paths
+    for ext in image_extensions:
+        # Use glob to find all image files with the given extension
+        image_paths.extend(glob.glob(os.path.join(save_dir, f'*.{ext}')))
+
+    # Now image_paths contains the list of all image file paths
+    print(image_paths)
+
+    return image_paths
 
 def generate_image(prompt, model, num_iterations, learning_rate, progress=gr.Progress(track_tqdm=True)):
     # Set up arguments
@@ -40,12 +58,13 @@ def generate_image(prompt, model, num_iterations, learning_rate, progress=gr.Pro
         image_path = f"{save_dir}/best_image.png"
         
         if os.path.exists(image_path):
-            return image_path, f"Image generated successfully and saved at {image_path}"
+            iter_images = list_iter_images(save_dir)
+            return image_path, f"Image generated successfully and saved at {image_path}", iter_images
         else:
-            return None, "Image generation completed, but the file was not found."
+            return None, "Image generation completed, but the file was not found.", None
     
     except Exception as e:
-        return None, f"An error occurred: {str(e)}"
+        return None, f"An error occurred: {str(e)}", None
 
 # Create Gradio interface
 title="# ReNO: Enhancing One-step Text-to-Image Models through Reward-based Noise Optimization"
@@ -92,11 +111,12 @@ with gr.Blocks() as demo:
             with gr.Column():
                 output_image = gr.Image(type="filepath", label="Generated Image")
                 status = gr.Textbox(label="Status")
+                iter_gallery = gr.Gallery(label="Iterations")
 
     submit_btn.click(
         fn = generate_image,
         inputs = [prompt, chosen_model, n_iter, learning_rate],
-        outputs = [output_image, status]
+        outputs = [output_image, status, iter_gallery]
     )
 
 # Launch the app