Update app.py

app.py

@@ -8,14 +8,16 @@ from diffusion_lens import get_images
 def generate_images(prompt):
     print('calling diffusion lens')
     all_images = []  # Initialize a list to store all images
-    for skip_layers in range(11, 0, -1):
+    all_labels = []
+    for skip_layers in range(12, -1, -1):
         images = get_images(prompt, skip_layers=skip_layers)
         all_images.append(images[0])  # Add the new image to the list
-        yield all_images  # Yield the list of all images
+        all_labels.append(f'layer_{12 - skip_layers}')
+        yield all_images, all_labels  # Yield the list of all images
 
 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")
+    gallery = gr.Gallery(label="Generated Images", columns=7, rows=2, object_fit="contain", height="auto")
     text_input.change(fn=generate_images, inputs=text_input, outputs=gallery)
 
 demo.launch()