From c7aede30469a9b0e2b47bb2e04bc042dfa6e2f4c Mon Sep 17 00:00:00 2001 From: Simon Lynch Date: Sun, 21 Jun 2026 21:25:13 +1000 Subject: [PATCH 1/2] Fix #8673: SD3 attention masks for text padding tokens --- src/diffusers/models/attention_processor.py | 45 ++++++++++++++++++- .../test_models_transformer_sd3.py | 43 ++++++++++++++++++ 2 files changed, 86 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index e2ece5cb3685..4751225face4 100755 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -754,6 +754,37 @@ def prepare_attention_mask( return attention_mask + def prepare_joint_attention_mask( + self, attention_mask: torch.Tensor, target_length: int, dtype: torch.dtype + ) -> torch.Tensor: + r""" + Prepare a joint attention mask for MMDiT-style processors that concatenate + `[hidden_states, encoder_hidden_states]` before attention (SD3 order). + + The input mask covers text tokens only; image/latent tokens are left-padded with + `1.0` (attend) so padding respects the concat order. + """ + if attention_mask is None: + return attention_mask + + current_length: int = attention_mask.shape[-1] + remaining_length: int = target_length - current_length + if current_length != target_length: + if attention_mask.device.type == "mps": + padding_shape = (attention_mask.shape[0], remaining_length) + padding = torch.ones(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device) + attention_mask = torch.cat([padding, attention_mask], dim=1) + else: + attention_mask = F.pad(attention_mask, (remaining_length, 0), value=1.0) + + if attention_mask.dim() == 3: + attention_mask = attention_mask[:, None, :, :] + elif attention_mask.dim() == 2: + attention_mask = attention_mask[:, None, None, :] + + attention_mask = (1.0 - attention_mask.to(dtype)) * torch.finfo(dtype).min + return attention_mask + def norm_encoder_hidden_states(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor: r""" Normalize the encoder hidden states. Requires `self.norm_cross` to be specified when constructing the @@ -1481,7 +1512,12 @@ def __call__( key = torch.cat([key, encoder_hidden_states_key_proj], dim=2) value = torch.cat([value, encoder_hidden_states_value_proj], dim=2) - hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) + if attention_mask is not None: + attention_mask = attn.prepare_joint_attention_mask(attention_mask, key.shape[2], key.dtype) + + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) hidden_states = hidden_states.to(query.dtype) @@ -1880,7 +1916,12 @@ def __call__( key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False) + if attention_mask is not None: + attention_mask = attn.prepare_joint_attention_mask(attention_mask, key.shape[2], key.dtype) + + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) hidden_states = hidden_states.to(query.dtype) diff --git a/tests/models/transformers/test_models_transformer_sd3.py b/tests/models/transformers/test_models_transformer_sd3.py index 50590fab8be1..66a730bc6fc2 100644 --- a/tests/models/transformers/test_models_transformer_sd3.py +++ b/tests/models/transformers/test_models_transformer_sd3.py @@ -224,3 +224,46 @@ class TestSD35TransformerBitsAndBytes(SD35TransformerTesterConfig, BitsAndBytesT class TestSD35TransformerTorchAo(SD35TransformerTesterConfig, TorchAoTesterMixin): """TorchAO quantization tests for SD3.5 Transformer.""" + + +class TestSD3JointAttentionMask(SD3TransformerTesterConfig): + def test_joint_attention_mask_makes_padding_invariant(self): + init_dict = self.get_init_dict() + model = SD3Transformer2DModel(**init_dict).to(torch_device) + model.eval() + + inputs = self.get_dummy_inputs(batch_size=1) + hidden_states = inputs["hidden_states"] + pooled_projections = inputs["pooled_projections"] + timestep = inputs["timestep"] + + content_length = 80 + padded_length = 154 + embedding_dim = init_dict["joint_attention_dim"] + + content = torch.randn(1, content_length, embedding_dim, generator=self.generator, device=torch_device) + padded = torch.zeros(1, padded_length, embedding_dim, device=torch_device, dtype=content.dtype) + padded[:, :content_length] = content + + attention_mask = torch.zeros(1, padded_length, device=torch_device) + attention_mask[:, :content_length] = 1.0 + + with torch.no_grad(): + out_padded = model( + hidden_states=hidden_states, + encoder_hidden_states=padded, + pooled_projections=pooled_projections, + timestep=timestep, + joint_attention_kwargs={"attention_mask": attention_mask}, + return_dict=False, + )[0] + + out_trimmed = model( + hidden_states=hidden_states, + encoder_hidden_states=content, + pooled_projections=pooled_projections, + timestep=timestep, + return_dict=False, + )[0] + + assert torch.allclose(out_padded, out_trimmed, atol=1e-5, rtol=1e-4) From f41465a2a572e223fd8d71c67ae4afdc0c4e79b2 Mon Sep 17 00:00:00 2001 From: Simon Lynch Date: Sun, 21 Jun 2026 21:46:16 +1000 Subject: [PATCH 2/2] Add processor-level SD3 joint attention mask test Standalone test avoids full transformer import chain; verifies padding invariance at processor level. Co-authored-by: Cursor --- tests/models/test_sd3_joint_attention_mask.py | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) create mode 100644 tests/models/test_sd3_joint_attention_mask.py diff --git a/tests/models/test_sd3_joint_attention_mask.py b/tests/models/test_sd3_joint_attention_mask.py new file mode 100644 index 000000000000..706116c2bbfb --- /dev/null +++ b/tests/models/test_sd3_joint_attention_mask.py @@ -0,0 +1,26 @@ +# coding=utf-8 +import torch +from diffusers.models.attention_processor import Attention, JointAttnProcessor2_0 +from ..testing_utils import enable_full_determinism, torch_device +enable_full_determinism() + +class TestSD3JointAttentionMaskProcessor: + def test_joint_attention_mask_makes_padding_invariant(self): + attn = Attention( + query_dim=32, added_kv_proj_dim=32, dim_head=8, heads=4, out_dim=32, + context_pre_only=False, bias=True, processor=JointAttnProcessor2_0(), eps=1e-6, + ).eval().to(torch_device) + batch_size, image_seq_len, short_text_len, long_text_len, dim = 1, 64, 20, 40, 32 + gen = torch.Generator(device=torch_device).manual_seed(0) + hidden_states = torch.randn(batch_size, image_seq_len, dim, generator=gen, device=torch_device) + encoder_short = torch.randn(batch_size, short_text_len, dim, generator=gen, device=torch_device) + encoder_long = torch.zeros(batch_size, long_text_len, dim, device=torch_device, dtype=encoder_short.dtype) + encoder_long[:, :short_text_len] = encoder_short + mask_short = torch.ones((batch_size, short_text_len), device=torch_device) + mask_long = torch.zeros((batch_size, long_text_len), device=torch_device) + mask_long[:, :short_text_len] = 1.0 + with torch.no_grad(): + out_s_hs, out_s_enc = attn(hidden_states=hidden_states, encoder_hidden_states=encoder_short, attention_mask=mask_short) + out_l_hs, out_l_enc = attn(hidden_states=hidden_states, encoder_hidden_states=encoder_long, attention_mask=mask_long) + assert (out_s_hs - out_l_hs).abs().max().item() < 1e-5 + assert (out_s_enc - out_l_enc[:, :short_text_len]).abs().max().item() < 1e-5