import spaces
import rembg
import torch
from diffusers import DiffusionPipeline

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