app updated

app.py

@@ -7,21 +7,44 @@ import gradio as gr
 import os
 import subprocess
 
-def get_prompt(prompt):
-    print('prompt:', prompt)
-    print('calling diffusion lens')
-    image = get_images(prompt, skip_layers=0)
-    return image
+
 
 def display_images(images):
     # Prepare images for display
     return [gr.Image(image) for image in images]
 
+# def get_prompt(prompt):
+#     print('prompt:', prompt)
+#     print('calling diffusion lens')
+#     image = get_images(prompt, skip_layers=0)
+#     return image
+#
+# if __name__ == '__main__':
+#     with gr.Blocks() as demo:
+#         gallery = gr.Gallery(
+#             label="Generated images", show_label=False, elem_id="gallery",
+#             columns=[1], rows=[1], object_fit="contain", height="auto")
+#         btn = gr.Button("Generate images", scale=0)
+#
+#         btn.click(get_prompt, 'text', gallery)
+#
+#     demo.launch()
+
+def get_prompt(prompt):
+    print('prompt:', prompt)
+    print('calling diffusion lens')
+
+    for skip_layers in range(24):  # loop from 0 to 23
+        image = get_images(prompt, skip_layers=skip_layers)
+        gallery.update(image)  # update the gallery with the new image
+
+    return "Images generated"
+
 if __name__ == '__main__':
     with gr.Blocks() as demo:
         gallery = gr.Gallery(
             label="Generated images", show_label=False, elem_id="gallery",
-            columns=[1], rows=[1], object_fit="contain", height="auto")
+            columns=[6], rows=[4], object_fit="contain", height="auto")  # set rows to 24 to accommodate all images
         btn = gr.Button("Generate images", scale=0)
 
         btn.click(get_prompt, 'text', gallery)