Update app.py

app.py

@@ -5,9 +5,9 @@ from diffusion_lens import get_images
 
 
 
-def generate_images(prompt):
+def generate_images(prompt, all_images):
     print('calling diffusion lens')
-    all_images = []  # Initialize a list to store all images
+    # all_images = []  # Initialize a list to store all images
     for skip_layers in range(11, -1, -1):
         images = get_images(prompt, skip_layers=skip_layers)
         all_images.append(images[0]) # (images[0], f'layer_{12 - skip_layers}'))  # Add the new image to the list
@@ -16,6 +16,13 @@ def generate_images(prompt):
 with gr.Blocks() as demo:
     text_input = gr.Textbox(label="Enter prompt")
     gallery = gr.Gallery(label="Generated Images", columns=6, rows=2, object_fit="contain", height="auto")
+
+    button = gr.Button("Diffusion Lens")  # Create a button with the label 'Diffusion Lens'
+    all_images = []  # Initialize a list to store all images outside the function
+
+    # Bind the button click to the generate_images function
+    button.click(fn=generate_images, inputs=[text_input, gr.State(all_images)], outputs=gallery)
+    
     text_input.change(fn=generate_images, inputs=text_input, outputs=gallery)
 
 demo.launch()