Update app.py

app.py

@@ -97,7 +97,7 @@ class ModelWrapper:
         
         for constant in all_timesteps:
             current_timesteps = torch.ones(len(prompt_embed), device="cuda", dtype=torch.long) * constant
-            current_timesteps = current_timesteps.to(torch.float16)
+            #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))
@@ -135,8 +135,8 @@ class ModelWrapper:
         prompt_embeds, pooled_prompt_embeds = self.text_encoder(prompt_inputs)
 
         batch_prompt_embeds, batch_pooled_prompt_embeds = (
-            prompt_embeds.repeat(num_images, 1, 1).to(torch.float16),
-            pooled_prompt_embeds.repeat(num_images, 1, 1).to(torch.float16)
+            prompt_embeds.repeat(num_images, 1, 1),
+            pooled_prompt_embeds.repeat(num_images, 1, 1)
         )
 
         unet_added_conditions = {