Trying fix for vid distortion

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -375,8 +375,6 @@ def __call__(self, x: jax.Array, feat_cache=None, feat_idx=0):
         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: