app.py

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():
        #  show images
            gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery", show_download_button=True, columns=[2])
    

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

ShoeGen.launch()