import gradio as gr
from diffusion_lens import get_images


def generate_images(prompt):
    print('calling diffusion lens')
    gr.Info('Generating images from intermediate layers..')
    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], f'layer_{12 - skip_layers}'))  
        yield 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")
    text_input.submit(fn=generate_images, inputs=text_input, outputs=gallery)

demo.launch()