import gradio as gr
import sys
# sys.path.append("LaVi-Bridge/test")
# from llama2_unet_diffusion_lens import call_diffusion_lens
from diffusion_lens import get_images
import gradio as gr
import os
import subprocess



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)
    print('calling diffusion lens')

    for skip_layers in range(24):  # loop from 0 to 23
        image = get_images(prompt, skip_layers=skip_layers)
        gallery.update(image)  # update the gallery with the new image

    return "Images generated"



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='Gallery')
    btn.click(get_prompt, text_input, gallery)  # pass the text input interface to btn.click()
demo.launch()