Preserve prompt padding masks in cross-attention

src/model/easywam_hidden.py

@@ -10,6 +10,7 @@
 from .component.action_dit import ActionDiT, StateEncoder
 from .helpers.gradient import gradient_checkpoint_forward
 from .backbone.wan22.loader import load_wan22_ti2v_5b_components
+from .prompt_utils import mask_prompt_padding
 from .schedulers.scheduler_continuous import ContinuousFlowMatchScheduler
 
 
@@ -539,10 +540,7 @@ def encode_prompt(self, prompt: Union[str, Sequence[str]]):
         ids = ids.to(self.device)
         mask = mask.to(self.device, dtype=torch.bool)
         prompt_emb = self.text_encoder(ids, mask)
-        seq_lens = mask.gt(0).sum(dim=1).long()
-        for i, v in enumerate(seq_lens):
-            prompt_emb[i, v:] = 0
-        mask = torch.ones_like(mask)
+        prompt_emb, mask = mask_prompt_padding(prompt_emb, mask)
         return prompt_emb.to(device=self.device), mask
 
     @torch.no_grad()

src/model/easywam_mot.py

@@ -10,6 +10,7 @@
 from .component.action_dit import ActionDiT, StateEncoder
 from .component.attention import AttentionSegment, StructuredAttentionMask, build_structured_attention_mask
 from .component.mot import MoT
+from .prompt_utils import mask_prompt_padding
 from .schedulers.scheduler_continuous import ContinuousFlowMatchScheduler
 
 logger = get_logger(__name__)
@@ -351,11 +352,7 @@ def encode_prompt(self, prompt: Union[str, Sequence[str]]):
         ids = ids.to(self.device)
         mask = mask.to(self.device, dtype=torch.bool)
         prompt_emb = self.text_encoder(ids, mask)
-        # FIXME: original implementation's zero padding is visible in cross-attn.
-        seq_lens = mask.gt(0).sum(dim=1).long()
-        for i, v in enumerate(seq_lens):
-            prompt_emb[i, v:] = 0
-        mask = torch.ones_like(mask)
+        prompt_emb, mask = mask_prompt_padding(prompt_emb, mask)
         return prompt_emb.to(device=self.device), mask
 
     def _append_state_to_context(

src/model/easywam_unified.py

@@ -14,6 +14,7 @@
     build_structured_attention_mask,
 )
 from .backbone.wan22.loader import load_wan22_ti2v_5b_components
+from .prompt_utils import mask_prompt_padding
 from .schedulers.scheduler_continuous import ContinuousFlowMatchScheduler
 from .schedulers.scheduler_flow_unipc import FlowUniPCScheduler
 
@@ -412,10 +413,7 @@ def encode_prompt(self, prompt: Union[str, Sequence[str]]):
         ids = ids.to(self.device)
         mask = mask.to(self.device, dtype=torch.bool)
         prompt_emb = self.text_encoder(ids, mask)
-        seq_lens = mask.gt(0).sum(dim=1).long()
-        for i, v in enumerate(seq_lens):
-            prompt_emb[i, v:] = 0
-        mask = torch.ones_like(mask)
+        prompt_emb, mask = mask_prompt_padding(prompt_emb, mask)
         return prompt_emb.to(device=self.device), mask
 
     @torch.no_grad()

src/model/prompt_utils.py

@@ -0,0 +1,37 @@
+"""Utilities for preparing text-conditioning tensors."""
+
+from __future__ import annotations
+
+import torch
+
+
+def mask_prompt_padding(
+    prompt_embeddings: torch.Tensor,
+    attention_mask: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Zero padded prompt rows while retaining their attention mask.
+
+    Text encoders return a fixed-width sequence for batched prompts.  Keeping
+    the tokenizer mask is important: a zero embedding is still a valid key to
+    attention unless its corresponding mask entry is false.
+    """
+
+    if prompt_embeddings.ndim != 3:
+        raise ValueError(
+            "prompt_embeddings must be [B, L, D], "
+            f"got {tuple(prompt_embeddings.shape)}"
+        )
+    if attention_mask.ndim != 2:
+        raise ValueError(
+            "attention_mask must be [B, L], "
+            f"got {tuple(attention_mask.shape)}"
+        )
+    if prompt_embeddings.shape[:2] != attention_mask.shape:
+        raise ValueError(
+            "prompt_embeddings and attention_mask must agree on [B, L], "
+            f"got {tuple(prompt_embeddings.shape[:2])} and {tuple(attention_mask.shape)}"
+        )
+
+    mask = attention_mask.to(device=prompt_embeddings.device, dtype=torch.bool)
+    prompt_embeddings = prompt_embeddings.masked_fill(~mask.unsqueeze(-1), 0)
+    return prompt_embeddings, mask

tests/test_prompt_utils.py

@@ -0,0 +1,34 @@
+"""CPU regression tests for text-conditioning padding."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from model.prompt_utils import mask_prompt_padding
+
+
+def test_mask_prompt_padding_hides_only_padded_tokens() -> None:
+    embeddings = torch.arange(24, dtype=torch.float32).reshape(2, 4, 3)
+    original = embeddings.clone()
+    attention_mask = torch.tensor([[1, 1, 0, 0], [1, 0, 1, 0]], dtype=torch.int64)
+
+    masked_embeddings, returned_mask = mask_prompt_padding(embeddings, attention_mask)
+
+    expected = original.masked_fill(~attention_mask.bool().unsqueeze(-1), 0)
+    torch.testing.assert_close(masked_embeddings, expected)
+    torch.testing.assert_close(returned_mask, attention_mask.bool())
+    assert masked_embeddings.dtype == embeddings.dtype
+    # The helper must not mutate a text encoder's output in-place.
+    torch.testing.assert_close(embeddings, original)
+
+
+@pytest.mark.parametrize(
+    ("embeddings_shape", "mask_shape"),
+    [((2, 4), (2, 4)), ((2, 4, 8), (4,)), ((2, 4, 8), (2, 3))],
+)
+def test_mask_prompt_padding_rejects_incompatible_shapes(
+    embeddings_shape: tuple[int, ...], mask_shape: tuple[int, ...]
+) -> None:
+    with pytest.raises(ValueError):
+        mask_prompt_padding(torch.zeros(embeddings_shape), torch.ones(mask_shape))