From d3983f9cebb85ab5773d13491d1b85ac3bf516f4 Mon Sep 17 00:00:00 2001 From: nan Date: Thu, 3 Sep 2026 00:42:44 +0800 Subject: [PATCH] Preserve prompt padding masks in cross-attention --- src/model/easywam_hidden.py | 6 ++---- src/model/easywam_mot.py | 7 ++----- src/model/easywam_unified.py | 6 ++---- src/model/prompt_utils.py | 37 ++++++++++++++++++++++++++++++++++++ tests/test_prompt_utils.py | 34 +++++++++++++++++++++++++++++++++ 5 files changed, 77 insertions(+), 13 deletions(-) create mode 100644 src/model/prompt_utils.py create mode 100644 tests/test_prompt_utils.py diff --git a/src/model/easywam_hidden.py b/src/model/easywam_hidden.py index 6c169c3..077878f 100644 --- a/src/model/easywam_hidden.py +++ b/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() diff --git a/src/model/easywam_mot.py b/src/model/easywam_mot.py index d4108d6..341fc52 100644 --- a/src/model/easywam_mot.py +++ b/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( diff --git a/src/model/easywam_unified.py b/src/model/easywam_unified.py index 648fe70..c7261d0 100644 --- a/src/model/easywam_unified.py +++ b/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() diff --git a/src/model/prompt_utils.py b/src/model/prompt_utils.py new file mode 100644 index 0000000..64ade19 --- /dev/null +++ b/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 diff --git a/tests/test_prompt_utils.py b/tests/test_prompt_utils.py new file mode 100644 index 0000000..890f4d0 --- /dev/null +++ b/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))