Update app.py

app.py

@@ -105,7 +105,7 @@ class ModelWrapper:
             eval_images = get_x0_from_noise(noise, eval_images, alphas_cumprod, current_timesteps).to(self.DTYPE)
             print(eval_images.dtype)
             next_timestep = current_timesteps - step_interval 
-            noise = self.scheduler.add_noise(eval_images, torch.randn_like(eval_images), next_timestep).to(torch.float16) 
+            noise = self.scheduler.add_noise(eval_images, torch.randn_like(eval_images), next_timestep).to(DTYPE) 
             print(noise.dtype)
         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]
@@ -210,6 +210,8 @@ def create_demo():
     conditioning_timestep = 999
     num_step = 4
     revision = None
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True 
 
     accelerator = Accelerator()