Update app.py

app.py

@@ -47,7 +47,6 @@ class ModelWrapper:
         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")
@@ -149,6 +148,7 @@ class ModelWrapper:
 
         return output_image_list, f"Run successfully in {(end_time-start_time):.2f} seconds"
 
+@spaces.GPU()
 def get_x0_from_noise(sample, model_output, alphas_cumprod, timestep):
     alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
     beta_prod_t = 1 - alpha_prod_t
@@ -185,6 +185,7 @@ class SDXLTextEncoder(torch.nn.Module):
         
         return prompt_embeds, pooled_prompt_embeds
 
+
 def create_demo():
     TITLE = "# DMD2-SDXL Demo"
     model_id = "stabilityai/stable-diffusion-xl-base-1.0"