import gradio as gr
from inference import generate_images

def infer(prompt):
    images = generate_images(prompt, model_path="output/", num_images=1)
    return images[0]

with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("## Image Generation Demo")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", lines=2)
            submit = gr.Button("Generate Image")
        with gr.Column():
            output_image = gr.Image(label="Generated Image")
    
    submit.click(infer, inputs=prompt, outputs=output_image)

if __name__ == "__main__":
    demo.launch()