Enable Flash Attention by trimming prompt embedding padding

pipeline_motif_video.py

@@ -539,6 +539,11 @@ class MotifVideoPipeline(DiffusionPipeline):
                 **prompt_embeds_kwargs,
             )
 
+        # Compute actual (non-padding) token count for batch=1 Flash Attention trimming in __call__
+        actual_seq_len = None
+        if batch_size == 1 and prompt_attention_mask is not None:
+            actual_seq_len = int(prompt_attention_mask.sum(dim=-1).max().item())
+
         # duplicate text embeddings for each generation per prompt, using mps friendly method
         seq_len = prompt_embeds.shape[1]
         prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
@@ -547,15 +552,16 @@ class MotifVideoPipeline(DiffusionPipeline):
         if pooled_prompt_embeds is not None:
             pooled_prompt_embeds = pooled_prompt_embeds.repeat_interleave(num_videos_per_prompt, dim=0)
 
-        # Keep attention mask handling
-        prompt_attention_mask = prompt_attention_mask.bool()
-        prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
-        prompt_attention_mask = prompt_attention_mask.repeat_interleave(num_videos_per_prompt, dim=0)
+        if prompt_attention_mask is not None:
+            prompt_attention_mask = prompt_attention_mask.bool()
+            prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
+            prompt_attention_mask = prompt_attention_mask.repeat_interleave(num_videos_per_prompt, dim=0)
 
         return (
             prompt_embeds,
             pooled_prompt_embeds,
             prompt_attention_mask,
+            actual_seq_len,
         )
 
     @property
@@ -1081,7 +1087,7 @@ class MotifVideoPipeline(DiffusionPipeline):
         device = self._execution_device
 
         # 3. Prepare text embeddings
-        prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self.encode_prompt(
+        prompt_embeds, pooled_prompt_embeds, prompt_attention_mask, pos_actual_len = self.encode_prompt(
             prompt=prompt,
             num_videos_per_prompt=num_videos_per_prompt,
             prompt_embeds=prompt_embeds,
@@ -1091,12 +1097,17 @@ class MotifVideoPipeline(DiffusionPipeline):
             device=device,
         )
 
+        if not self.guider._enabled and pos_actual_len is not None:
+            prompt_embeds = prompt_embeds[:, :pos_actual_len, :]
+            prompt_attention_mask = None
+
         if self.guider._enabled:
             negative_prompt = self._prepare_negative_prompt(negative_prompt, batch_size)
             (
                 negative_prompt_embeds,
                 negative_pooled_prompt_embeds,
                 negative_prompt_attention_mask,
+                neg_actual_len,
             ) = self.encode_prompt(
                 prompt=negative_prompt,
                 num_videos_per_prompt=num_videos_per_prompt,
@@ -1107,6 +1118,14 @@ class MotifVideoPipeline(DiffusionPipeline):
                 device=device,
             )
 
+            # Trim each to its own actual length — guider runs pos/neg in separate loop iterations,
+            # so different seq lengths are fine. No padding embeddings attend without mask.
+            if pos_actual_len is not None and neg_actual_len is not None:
+                prompt_embeds = prompt_embeds[:, :pos_actual_len, :]
+                negative_prompt_embeds = negative_prompt_embeds[:, :neg_actual_len, :]
+                prompt_attention_mask = None
+                negative_prompt_attention_mask = None
+
         num_channels_latents = self.vae.config.z_dim
         latents = self.prepare_latents(
             batch_size * num_videos_per_prompt,
@@ -1229,7 +1248,7 @@ class MotifVideoPipeline(DiffusionPipeline):
                 guider_inputs = {
                     "encoder_hidden_states": (prompt_embeds, negative_prompt_embeds),
                 }
-                if use_attention_mask:
+                if use_attention_mask and prompt_attention_mask is not None:
                     guider_inputs["encoder_attention_mask"] = (
                         prompt_attention_mask,
                         negative_prompt_attention_mask,