image_embeds sharded

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -637,14 +637,12 @@ 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())
-    print(f"[DEBUG PREP] prompt_sharding spec: {self.config.data_sharding}")
-    print(f"[DEBUG PREP] image_sharding spec: () - Replicated")
-
-    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)
+    data_sharding = NamedSharding(self.mesh, P(*self.config.data_sharding))
+    print(f"[DEBUG PREP] data_sharding spec: {self.config.data_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)
 
     print(f"[DEBUG PREP] SHARDED prompt_embeds.shape: {prompt_embeds.shape}")
     print(f"[DEBUG PREP] SHARDED image_embeds.shape: {image_embeds.shape}")