Update app.py

app.py

@@ -92,6 +92,7 @@ def scheduler_swap_callback(pipeline, step_index, timestep, callback_kwargs):
         torch.backends.cuda.preferred_blas_library="cublaslt"
     if step_index == int(pipeline.num_timesteps * 0.5):
        # torch.set_float32_matmul_precision("medium")
+        callback_kwargs["latents"].to(torch.float64)
         pipe.unet.to(torch.float64)
       #  pipe.guidance_scale=1.0
        # pipe.scheduler.set_timesteps(num_inference_steps*.70)
@@ -102,6 +103,7 @@ def scheduler_swap_callback(pipeline, step_index, timestep, callback_kwargs):
         torch.backends.cudnn.allow_tf32 = False
         torch.backends.cuda.matmul.allow_tf32 = False
         torch.set_float32_matmul_precision("highest")
+        callback_kwargs["latents"].to(torch.bfloat16)
         pipe.unet.to(torch.bfloat16)
      #   pipe.vae = vae_a
      #   pipe.unet = unet_a