Update app.py

app.py

@@ -96,6 +96,7 @@ class main():
     @torch.no_grad()
     @spaces.GPU
     def inference(self, prompt, negative_prompt, guidance_scale, ddim_steps, seed):
+        device = self.device
         generator = torch.Generator(device=device).manual_seed(seed)
         latents = torch.randn(
                 (1, self.unet.in_channels, 512 // 8, 512 // 8),
@@ -140,6 +141,7 @@ class main():
     @torch.no_grad()
     @spaces.GPU
     def edit_inference(self, prompt, negative_prompt, guidance_scale, ddim_steps, seed, start_noise, a1, a2, a3, a4): 
+        device = self.device
         original_weights = self,network.proj.clone()
             
         #pad to same number of PCs