Trying text_mask 9

Author: prishajain1

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -633,7 +633,7 @@ def _prepare_model_inputs_i2v(
     effective_batch_size = batch_size * num_videos_per_prompt
 
     # 1. Encode Prompts
-    prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+    prompt_embeds, negative_prompt_embeds, text_attention_mask, negative_text_attention_mask = self.encode_prompt(
         prompt=prompt,
         negative_prompt=negative_prompt,
         num_videos_per_prompt=num_videos_per_prompt,
@@ -677,9 +677,12 @@ def _prepare_model_inputs_i2v(
 
     prompt_embeds = jax.device_put(prompt_embeds, data_sharding)
     negative_prompt_embeds = jax.device_put(negative_prompt_embeds, data_sharding)
+    text_attention_mask = jax.device_put(text_attention_mask, data_sharding)
+    negative_text_attention_mask = jax.device_put(negative_text_attention_mask, data_sharding)
     image_embeds = jax.device_put(image_embeds, data_sharding)
 
-    return prompt_embeds, negative_prompt_embeds, image_embeds, effective_batch_size
+    return (prompt_embeds, negative_prompt_embeds, text_attention_mask, 
+            negative_text_attention_mask, image_embeds, effective_batch_size)
 
 
   def _prepare_model_inputs(