cc

scripts/compute_latent.py

@@ -4,6 +4,7 @@ import PIL
 
 from itertools import batched
 
+import numpy as np
 import torch
 import torchvision.transforms as T
 from diffusers import AutoencoderKL
@@ -51,11 +52,11 @@ def compute_latent(model_name: str, swim_dir: str, batch_size: int):
                 latents = latents.detach().cpu().numpy()
 
             for name, latent in zip(image_names, latents):
-                torch.save(
-                    latent,
+                np.save(
                     os.path.join(
-                        output, name.replace(".jpg", ".pt").replace(".png", ".pt")
+                        output, name.replace(".jpg", ".npy").replace(".png", ".npy")
                     ),
+                    latent,
                 )