import spaces
import gradio as gr
import torch
from diffusers import DiffusionPipeline

model_id = "stabilityai/stable-diffusion-2-1"
pipe = DiffusionPipeline.from_pretrained(model_id)

pipe.to("cuda")

# Function to generate an image from text using diffusion
@spaces.GPU
def generate_image(prompt):
    
    images = pipe(prompt).images
    return images

_TITLE = "Shoe Generator"
with gr.Blocks(_TITLE) as ShoeGen:
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Enter a prompt")
            button_gen = gr.Button("Generate Image")
        with gr.Column():
        #  show images
            gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[2], height="auto")
    

    button_gen.click(generate_image, inputs=[prompt], outputs=gallery)

ShoeGen.launch()