Update app.py

app.py

@@ -89,12 +89,12 @@ class ModelWrapper:
         else:
             raise NotImplementedError()
         
-        prompt_embed = prompt_embed.to(dtype=torch.float16)
         DTYPE = prompt_embed.dtype
         print(DTYPE)
         
         for constant in all_timesteps:
             current_timesteps = torch.ones(len(prompt_embed), device="cuda", dtype=torch.long) * constant
+            print(current_timesteps.dtype)
             eval_images = self.model(noise, current_timesteps, prompt_embed, added_cond_kwargs=unet_added_conditions).sample
             print(type(eval_images))
 
@@ -122,7 +122,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="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.float32) 
 
         prompt_inputs = self._encode_prompt(prompt)