Update app.py

app.py

@@ -90,7 +90,7 @@ class ModelWrapper:
         else:
             raise NotImplementedError()
 
-        #noise = noise.to(torch.float16)
+        noise = noise.to(torch.float16)
         print(f'noise: {noise.dtype}')
         #prompt_embed = prompt_embed.to(torch.float32)
         DTYPE = prompt_embed.dtype
@@ -100,7 +100,7 @@ class ModelWrapper:
             current_timesteps = torch.ones(len(prompt_embed), device="cuda", dtype=torch.long) * constant
             #current_timesteps = current_timesteps.to(torch.float32)
             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, current_timesteps, prompt_embed, added_cond_kwargs=unet_added_conditions)
             print(eval_images.dtype)
             eval_images = get_x0_from_noise(noise, eval_images, alphas_cumprod, current_timesteps).to(self.DTYPE)
             print(eval_images.dtype)
@@ -140,7 +140,7 @@ class ModelWrapper:
         )
 
         unet_added_conditions = {
-            "time_ids": add_time_ids.to(torch.float32),
+            "time_ids": add_time_ids,
             "text_embeds": batch_pooled_prompt_embeds.squeeze(1)
         }