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

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):
    
    image = pipe(prompt).images
    image = rembg.remove(image)
    return image

_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():
            image = gr.Image(label="Generated Image", show_download_button=True)    

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

ShoeGen.launch()