app.py

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