Update app.py

app.py

@@ -46,7 +46,8 @@ class ModelWrapper:
         self.scheduler = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler")
         self.alphas_cumprod = self.scheduler.alphas_cumprod.to(self.device)
         self.num_step = num_step
-        
+    
+    @spaces.GPU()  
     def create_generator(self, model_id, checkpoint_path):
         generator = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet").to(self.DTYPE)
         state_dict = torch.load(checkpoint_path, map_location="cuda")
@@ -108,7 +109,7 @@ class ModelWrapper:
         eval_images = ((eval_images + 1.0) * 127.5).clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1)
         return eval_images 
         
-    @spaces.GPU()
+    
     @torch.no_grad()
     def inference(self, prompt, seed, height, width, num_images, fast_vae_decode):
         print("Running model inference...")