Update app.py

app.py

@@ -90,8 +90,9 @@ class ModelWrapper:
         else:
             raise NotImplementedError()
 
-        noise = noise.to(device="cuda", dtype=torch.float32)
+        noise = noise.to(device="cuda", dtype=torch.float16)
         print(f'noise: {noise.dtype}')
+        prompt_embed = prompt_embed.to(device="cuda", dtype=torch.float16)
         DTYPE = prompt_embed.dtype
         print(f'prompt_embed: {DTYPE}')
         
@@ -100,7 +101,6 @@ class ModelWrapper:
             #current_timesteps = current_timesteps.to(torch.float16)
             print(f'current_timestpes: {current_timesteps.dtype}')
             eval_images = self.model(noise, current_timesteps, prompt_embed, added_cond_kwargs=unet_added_conditions).sample
-            print(type(eval_images))
 
             eval_images = get_x0_from_noise(noise, eval_images, alphas_cumprod, current_timesteps).to(self.DTYPE)