Update app.py

app.py

@@ -76,7 +76,8 @@ class ModelWrapper:
     @staticmethod
     def _get_time():
         return time.time()
-
+        
+    @spaces.GPU()
     def sample(self, noise, unet_added_conditions, prompt_embed, fast_vae_decode):
         alphas_cumprod = self.scheduler.alphas_cumprod.to(self.device)
 
@@ -123,10 +124,8 @@ class ModelWrapper:
 
         add_time_ids = self.build_condition_input(height, width).repeat(num_images, 1)
 
-        noise = torch.randn(num_images, 4, height // self.vae_downsample_ratio, width // self.vae_downsample_ratio, generator=generator).to(device="cuda",dtype=torch.float16)
+        noise = torch.randn(num_images, 4, height // self.vae_downsample_ratio, width // self.vae_downsample_ratio, generator=generator).to(device="cuda", dtype=torch.float16)
         
-        noise = noise.to(torch.float16)
-
         prompt_inputs = self._encode_prompt(prompt)
         
         start_time = self._get_time()
@@ -159,7 +158,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