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 get_prompt(prompt):
    print('prompt:', prompt)
    print('calling diffusion lens')
    image = get_images(prompt, skip_layers=0)
    return image

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

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