sharding for image and prompt embeds made same

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -632,13 +632,11 @@ def _prepare_model_inputs_i2v(
     if negative_prompt_embeds is not None:
       negative_prompt_embeds = negative_prompt_embeds.astype(transformer_dtype)
 
-    prompt_sharding = NamedSharding(self.mesh, P(*self.config.data_sharding))
-    image_sharding = NamedSharding(self.mesh, P())
+    data_sharding = NamedSharding(self.mesh, P(*self.config.data_sharding))
 
-
-    prompt_embeds = jax.device_put(prompt_embeds, prompt_sharding)
-    negative_prompt_embeds = jax.device_put(negative_prompt_embeds, prompt_sharding)
-    image_embeds = jax.device_put(image_embeds, image_sharding)
+    prompt_embeds = jax.device_put(prompt_embeds, data_sharding)
+    negative_prompt_embeds = jax.device_put(negative_prompt_embeds, data_sharding)
+    image_embeds = jax.device_put(image_embeds, data_sharding)
 
     return prompt_embeds, negative_prompt_embeds, image_embeds, effective_batch_size