Trying fix for vid distortion

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -369,11 +369,16 @@ 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)
         else:
           cache_x = jnp.copy(x[:, -1:, :, :, :])
           x = self.time_conv(jnp.concatenate([feat_cache[idx][:, -1:, :, :, :], x], axis=1))
+          # Discard the first frame of output as it corresponds to the cached frame (already processed)
+          x = x[:, 1:, ...]
           feat_cache = _update_cache(feat_cache, idx, cache_x)
           feat_idx += 1
+      else:
+         x = self.time_conv(x)
 
     return x, feat_cache, feat_idx