Update app.py

app.py

@@ -15,9 +15,11 @@ import torch
 
 # Load Stable Diffusion Model
 def load_stable_diffusion_model():
-    pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float32)
-    pipe = pipe.to("cpu") 
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float32).to(device)
+    return pipe
 
+# Initialize the model globally
 sd_model = load_stable_diffusion_model()
 
 def remove_background_stable_diffusion(image_path):