Update app.py

app.py

@@ -31,6 +31,7 @@ class StableDiffusionLoader:
         """
         self.prompt = prompt
         self.pretrain_pipe = pretrain_pipe
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
         assert isinstance(self.prompt, str), 'Please enter a string into the prompt field'
         assert isinstance(self.pretrain_pipe, str), 'Please use value such as `CompVis/stable-diffusion-v1-4` for pretrained pipeline'