Update app.py

app.py

@@ -14,8 +14,11 @@ import json
 import torch
 
 # Load Stable Diffusion Model
-pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16)
-pipe = pipe.to("cuda")  # Change to "cpu" if you don't have a GPU
+def load_stable_diffusion_model():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    return StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float32).to(device)
+
+sd_model = load_stable_diffusion_model()
 
 def remove_background_stable_diffusion(image_path):
     image = Image.open(image_path).convert("RGB")