import gradio as gr
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, -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 after each new image is added

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'
    
    # Use the button to trigger the generate_images function
    button.click(fn=generate_images, inputs=[text_input, gr.State(None)], outputs=gallery)

demo.launch()


# import gradio as gr
# from diffusion_lens import get_images





# def generate_images(prompt, all_images):
#     print('calling diffusion lens')
#     # 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
#         yield all_images  # 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")

#     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()


# # def display_images(images):
# #     # Prepare images for display
# #     return [gr.Image(image) for image in images]

# # def get_prompt(prompt):
# #     print('prompt:', prompt)
# #     return prompt

# # def generate_images(prompt):
# #     print('calling diffusion lens')
# #     for skip_layers in range(23, 0, -1):
# #         images = get_images(prompt, skip_layers=skip_layers)
# #         yield images[0]  # Yield each image as soon as it's ready
# #         # yield gr.Image(images[0])  # Yield each image as soon as it's ready

# # with gr.Blocks() as demo:
# #     text_input = gr.Interface(fn=generate_images, inputs="text", outputs="image")
# # demo.launch()