Upload anytext.py

anytext.py

@@ -822,7 +822,10 @@ class AuxiliaryLatentModule(nn.Module):
         # get masked_x
         masked_img = ((edit_image.astype(np.float32) / 127.5) - 1.0) * (1 - np_hint)
         masked_img = np.transpose(masked_img, (2, 0, 1))
-        masked_img = torch.from_numpy(masked_img.copy()).float().to(self.device)
+        print("vae device", next(self.vae.parameters()).device)
+        print("masked_img device", self.device)
+        device = next(self.vae.parameters()).device
+        masked_img = torch.from_numpy(masked_img.copy()).float().to(device)
         if self.use_fp16:
             masked_img = masked_img.half()
         masked_x = (retrieve_latents(self.vae.encode(masked_img[None, ...])) * self.vae.config.scaling_factor).detach()