Added Debug to check vae encode and decode time

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -15,7 +15,7 @@
 """
 
 from typing import Tuple, List, Sequence, Union, Optional
-
+import time
 import flax
 import jax
 import jax.numpy as jnp
@@ -1096,7 +1096,10 @@ def encode(
       self, x: jax.Array, feat_cache: AutoencoderKLWanCache, return_dict: bool = True
   ) -> Union[FlaxAutoencoderKLOutput, Tuple[FlaxDiagonalGaussianDistribution]]:
     """Encode video into latent distribution."""
+    s0 = time.perf_counter()
     h = self._encode(x, feat_cache)
+    h.block_until_ready()
+    print(f"VAE Encode time: {time.perf_counter() - s0:.4f}s")
     posterior = FlaxDiagonalGaussianDistribution(h)
     if not return_dict:
       return (posterior,)
@@ -1145,7 +1148,10 @@ def decode(
       # reshape channel last for JAX
       z = jnp.transpose(z, (0, 2, 3, 4, 1))
       assert z.shape[-1] == self.z_dim, f"Expected input shape (N, D, H, W, {self.z_dim}, got {z.shape}"
+    s0 = time.perf_counter()
     decoded = self._decode(z, feat_cache).sample
+    decoded.block_until_ready()
+    print(f"VAE Decode time: {time.perf_counter() - s0:.4f}s")
     if not return_dict:
       return (decoded,)
     return FlaxDecoderOutput(sample=decoded)