sharding before call to encoder

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -134,11 +134,6 @@ def __call__(self, x: jax.Array, cache_x: Optional[jax.Array] = None, idx=-1) ->
     else:
       x_padded = x
 
-    if self.mesh is not None:
-      # (B, D, H, W, C)
-      if x_padded.shape[0] % self.mesh.shape['data'] == 0:
-        x_padded = with_sharding_constraint(x_padded, PartitionSpec('data', None, None, None, None))
-
     out = self.conv(x_padded)
     return out
 

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -529,6 +529,11 @@ def prepare_latents_i2v_base(
 
       vae_dtype = getattr(self.vae, "dtype", jnp.float32)
       video_condition = video_condition.astype(vae_dtype)
+      sharding_spec = P(self.config.mesh_axes[0], None, None, None, None) # typically P('data', None, ...)
+      video_condition = jax.lax.with_sharding_constraint(
+          video_condition, 
+          sharding_spec
+      )
 
       with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
           encoded_output = self.vae.encode(video_condition, self.vae_cache)[0].mode()