Update app.py

app.py

@@ -13,23 +13,6 @@ 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)
     return prompt
@@ -38,19 +21,20 @@ def get_prompt(prompt):
 
 def generate_images(prompt):
     print('calling diffusion lens')
-    for skip_layers in range(24):  # loop from 0 to 23
+    images = []
+    for skip_layers in range(1):  # loop from 0 to 23
         image = get_images(prompt, skip_layers=skip_layers)
-        gallery.update(image)  # update the gallery with the new image
+        images.append(image)
 
-    return "Images generated"
+    return images
 
 
 
 with gr.Blocks() as demo:
-    gallery = gr.Gallery(
-        label="Generated images", show_label=False, elem_id="gallery",
-        columns=[6], rows=[4], object_fit="contain", height="auto")  # set rows to 24 to accommodate all images
-    btn = gr.Button("Generate images", scale=0)
-    text_input = gr.Interface(fn=get_prompt, inputs="text", outputs="text")
-    btn.click(generate_images, text_input, gallery)  # pass the text input interface to btn.click()
-demo.launch()
+    # gallery = gr.Gallery(
+    #     label="Generated images", show_label=False, elem_id="gallery",
+    #     columns=[6], rows=[4], object_fit="contain", height="auto")  # set rows to 24 to accommodate all images
+    # btn = gr.Button("Generate images", scale=0)
+    text_input = gr.Interface(fn=get_prompt, inputs="text", outputs=["image"] * 1)
+    # btn.click(generate_images, text_input, gallery)  # pass the text input interface to btn.click()
+demo.launch()