Trying fix for vid distortion

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -369,7 +369,9 @@ def __call__(self, x: jax.Array, feat_cache=None, feat_idx=0):
         if feat_cache[idx] is None:
           feat_cache = _update_cache(feat_cache, idx, jnp.copy(x))
           feat_idx += 1
-          x = self.time_conv(x)
+          # Pad with 2 zeros to satisfy kernel size 3 and stride 2 for a single frame
+          x_padded = jnp.pad(x, ((0, 0), (2, 0), (0, 0), (0, 0), (0, 0)))
+          x = self.time_conv(x_padded)
         else:
           cache_x = jnp.copy(x[:, -1:, :, :, :])
           x = self.time_conv(jnp.concatenate([feat_cache[idx][:, -1:, :, :, :], x], axis=1))
@@ -382,7 +384,6 @@ def __call__(self, x: jax.Array, feat_cache=None, feat_idx=0):
 
     return x, feat_cache, feat_idx
 
-
 class WanResidualBlock(nnx.Module):
 
   def __init__(