Update app.py

app.py

@@ -96,13 +96,13 @@ class ModelWrapper:
             current_timesteps = torch.ones(len(prompt_embed), device="cuda", dtype=torch.long) * constant
             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
+            eval_images = self.model(noise.to(torch.float16), current_timesteps, prompt_embed.to(torch.float16), 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)
 
             next_timestep = current_timesteps - step_interval 
-            noise = self.scheduler.add_noise(eval_images, torch.randn_like(eval_images), next_timestep).to(DTYPE)  
+            noise = self.scheduler.add_noise(eval_images, torch.randn_like(eval_images), next_timestep).to(torch.float16) 
 
         if fast_vae_decode:
             eval_images = self.tiny_vae.decode(eval_images.to(self.tiny_vae_dtype) / self.tiny_vae.config.scaling_factor, return_dict=False)[0]
@@ -138,7 +138,7 @@ class ModelWrapper:
 
         unet_added_conditions = {
             "time_ids": add_time_ids,
-            "text_embeds": batch_pooled_prompt_embeds.squeeze(1)
+            "text_embeds": batch_pooled_prompt_embeds.squeeze(1).to(torch.float16)
         }