Merge pull request #11 from LightricksResearch/rm-dist-util

xora/models/autoencoders/vae_encode.py

@@ -6,8 +6,10 @@ from torch import Tensor
 
 from xora.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
 from xora.models.autoencoders.video_autoencoder import Downsample3D, VideoAutoencoder
-import xora.utils.dist_util
-
+try:
+    import torch_xla.core.xla_model as xm
+except:
+    pass
 
 def vae_encode(media_items: Tensor, vae: AutoencoderKL, split_size: int = 1, vae_per_channel_normalize=False) -> Tensor:
     """
@@ -54,10 +56,12 @@ def vae_encode(media_items: Tensor, vae: AutoencoderKL, split_size: int = 1, vae
         encode_bs = len(media_items) // split_size
         # latents = [vae.encode(image_batch).latent_dist.sample() for image_batch in media_items.split(encode_bs)]
         latents = []
-        dist_util.execute_graph()
+        if media_items.device.type == "xla":
+            xm.mark_step()
         for image_batch in media_items.split(encode_bs):
             latents.append(vae.encode(image_batch).latent_dist.sample())
-            dist_util.execute_graph()
+            if media_items.device.type == "xla":
+                xm.mark_step()
         latents = torch.cat(latents, dim=0)
     else:
         latents = vae.encode(media_items).latent_dist.sample()

xora/utils/dist_util.py

@@ -1,5 +0,0 @@
-from enum import Enum
-
-def execute_graph() -> None:
-    if _acceleration_type == AccelerationType.TPU:
-        xm.mark_step()