update

tokenflow_pnp.py

@@ -21,14 +21,6 @@ logging.set_verbosity_error()
 
 VAE_BATCH_SIZE = 10
 
-if torch.cuda.is_available():
-    device = "cuda"
-elif torch.backends.mps.is_available():
-    device = "mps"
-else:
-    device = "cpu"
-to = torch.float16 if device == 'cuda' else torch.float32
-
 class TokenFlow(nn.Module):
     def __init__(self, config, 
                  pipe,
@@ -275,7 +267,7 @@ class TokenFlow(nn.Module):
         denoised_latent = self.scheduler.step(noise_pred, t, x)['prev_sample']
         return denoised_latent
     
-    @torch.autocast(dtype=to, device_type=device)
+    @torch.autocast(dtype=torch.float16, device_type='cuda')
     def batched_denoise_step(self, x, t, indices):
         batch_size = self.config["batch_size"]
         denoised_latents = []