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

model_id = "CiroN2022/shoes"
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)
    return images[0]

_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():
            output = gr.Image(label="Generated Image")

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

ShoeGen.launch()