Update app.py

app.py

@@ -21,7 +21,7 @@ class ModelWrapper:
         torch.set_grad_enabled(False)
         
         self.DTYPE = torch.float16
-        self.device = accelerator.device
+        self.device = 0
 
         self.tokenizer_one = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", revision=revision, use_fast=False)
         self.tokenizer_two = AutoTokenizer.from_pretrained(model_id, subfolder="tokenizer", revision=revision, use_fast=False)
@@ -60,7 +60,7 @@ class ModelWrapper:
         crop_top_left = (0, 0)
 
         add_time_ids = list(original_size + crop_top_left + target_size)
-        add_time_ids = torch.tensor([add_time_ids], device=self.device, dtype=self.DTYPE)
+        add_time_ids = torch.tensor([add_time_ids], device="cuda", dtype=self.DTYPE)
         return add_time_ids
 
     def _encode_prompt(self, prompt):
@@ -92,7 +92,7 @@ class ModelWrapper:
         DTYPE = prompt_embed.dtype
         
         for constant in all_timesteps:
-            current_timesteps = torch.ones(len(prompt_embed), device=self.device, dtype=torch.long) * constant
+            current_timesteps = torch.ones(len(prompt_embed), device="cuda", dtype=torch.long) * constant
             eval_images = self.model(noise, current_timesteps, prompt_embed, added_cond_kwargs=unet_added_conditions).sample
 
             eval_images = get_x0_from_noise(noise, eval_images, alphas_cumprod, current_timesteps).to(self.DTYPE)
@@ -120,7 +120,7 @@ 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=self.device, dtype=self.DTYPE) 
+        noise = torch.randn(num_images, 4, height // self.vae_downsample_ratio, width // self.vae_downsample_ratio, generator=generator).to(device="cuda", dtype=self.DTYPE) 
 
         prompt_inputs = self._encode_prompt(prompt)
         
@@ -148,7 +148,6 @@ 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