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(23, 0, -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

with gr.Blocks() as demo:
    text_input = gr.Textbox(label="Enter prompt")
    gallery = gr.Gallery(label="Generated Images", columns=4, rows=6)
    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()