app.py

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

pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
pipe.to("cuda")

# Function to generate an image from text using diffusion
@spaces.GPU
def generate_image(prompt):
    prompt += "no background, side view, minimalist shot"
    
    image = pipe(prompt).images[0]
    image2 = rembg.remove(image)

    return image, image2

_TITLE = "Shoe Generator"
with gr.Blocks(_TITLE) as ShoeGen:
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Enter a discription of a shoe")
            # neg_prompt = gr.Textbox(label="Enter a negative prompt", value="low quality, watermark, ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, body out of frame, blurry, bad anatomy, blurred, watermark, grainy, signature, cut off, draft, closed eyes, text, logo")
            button_gen = gr.Button("Generate Image")
        with gr.Column():
            image = gr.Image(label="Generated Image", show_download_button=True) 
            image_nobg = gr.Image(label="Generated Image (No Background)", show_download_button=True) 
    
    button_gen.click(generate_image, inputs=[prompt], outputs=[image, image_nobg])

ShoeGen.launch()