Update app.py

app.py

@@ -80,7 +80,7 @@ class ModelWrapper:
     @spaces.GPU()
     def sample(self, noise, unet_added_conditions, prompt_embed, fast_vae_decode):
         alphas_cumprod = self.scheduler.alphas_cumprod.to(self.device)
-
+        print(alphas_cumprod)
         if self.num_step == 1:
             all_timesteps = [self.conditioning_timestep]
             step_interval = 0 
@@ -90,9 +90,9 @@ class ModelWrapper:
         else:
             raise NotImplementedError()
 
-        noise = noise.to(device="cuda", dtype=torch.float16)
+        noise = noise.to(torch.float16)
         print(f'noise: {noise.dtype}')
-        prompt_embed = prompt_embed.to(device="cuda", dtype=torch.float16)
+        prompt_embed = prompt_embed.to(torch.float16)
         DTYPE = prompt_embed.dtype
         print(f'prompt_embed: {DTYPE}')
         
@@ -145,10 +145,11 @@ class ModelWrapper:
         }
 
 
-        
+        print(f'noise: {noise.dtype}')
         print(f'prompt: {batch_prompt_embeds.dtype}')
         print(unet_added_conditions['time_ids'].dtype)
         print(unet_added_conditions['text_embeds'].dtype)
+        print("________")
         
         eval_images = self.sample(noise=noise, unet_added_conditions=unet_added_conditions, prompt_embed=batch_prompt_embeds, fast_vae_decode=fast_vae_decode)