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)
    return prompt



def generate_images(prompt):
    print('calling diffusion lens')
    all_images = []
    for skip_layers in range(2):
        images = get_images(prompt, skip_layers=skip_layers)
        all_images.append(images[0])
    print('after call images')
    return all_images



with gr.Blocks() as demo:
    text_input = gr.Interface(fn=generate_images, inputs="text", outputs=["image"]*2)
demo.launch()