Update app.py

app.py

@@ -57,7 +57,6 @@ def load_models(device):
     unet = UNet2DConditionModel.from_pretrained(
             pretrained_model_name_or_path, subfolder="unet", revision=revision
         )
-    
     unet.requires_grad_(False)
     unet.to(device, dtype=weight_dtype)
     vae.requires_grad_(False)
@@ -124,7 +123,7 @@ class main():
         self.vae.to(device, dtype=weight_dtype)
         self.text_encoder.to(device, dtype=weight_dtype)
         print("")
-        print(self.text_encoder.device)
+
         
         self.network = None
     
@@ -171,7 +170,8 @@ class main():
         self.thick = thick
 
 
-
+    @torch.no_grad()
+    @spaces.GPU
     def sample_model(self):
         self.unet, _, _, _, _ = load_models(self.device)
         self.network = sample_weights(self.unet, self.proj, self.mean, self.std, self.v[:, :1000], self.device, factor = 1.00)
@@ -181,6 +181,13 @@ class main():
     @spaces.GPU
     def inference(self, prompt, negative_prompt, guidance_scale, ddim_steps, seed):
         device = self.device
+        self.unet.to(device)
+        self.text_encoder.to(device)
+        self.vae.to(device)
+        self.tokenizer.to(device)
+        
+
+        
         generator = torch.Generator(device=device).manual_seed(seed)
         latents = torch.randn(
                 (1, self.unet.in_channels, 512 // 8, 512 // 8),