Fix QwenImage txt_seq_lens handling (#12702) · NotSoBot/nsb.api.diffusers@dad5cb5 · GitHub
Skip to content

Commit dad5cb5

Browse files
kashifsayakpaulyiyixuxucdutr
authored
Fix QwenImage txt_seq_lens handling (huggingface#12702)
* Fix QwenImage txt_seq_lens handling * formatting * formatting * remove txt_seq_lens and use bool mask * use compute_text_seq_len_from_mask * add seq_lens to dispatch_attention_fn * use joint_seq_lens * remove unused index_block * WIP: Remove seq_lens parameter and use mask-based approach - Remove seq_lens parameter from dispatch_attention_fn - Update varlen backends to extract seqlens from masks - Update QwenImage to pass 2D joint_attention_mask - Fix native backend to handle 2D boolean masks - Fix sage_varlen seqlens_q to match seqlens_k for self-attention Note: sage_varlen still producing black images, needs further investigation * fix formatting * undo sage changes * xformers support * hub fix * fix torch compile issues * fix tests * use _prepare_attn_mask_native * proper deprecation notice * add deprecate to txt_seq_lens * Update src/diffusers/models/transformers/transformer_qwenimage.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * Update src/diffusers/models/transformers/transformer_qwenimage.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * Only create the mask if there's actual padding * fix order of docstrings * Adds performance benchmarks and optimization details for QwenImage Enhances documentation with comprehensive performance insights for QwenImage pipeline: * rope_text_seq_len = text_seq_len * rename to max_txt_seq_len * removed deprecated args * undo unrelated change * Updates QwenImage performance documentation Removes detailed attention backend benchmarks and simplifies torch.compile performance description Focuses on key performance improvement with torch.compile, highlighting the specific speedup from 4.70s to 1.93s on an A100 GPU Streamlines the documentation to provide more concise and actionable performance insights * Updates deprecation warnings for txt_seq_lens parameter Extends deprecation timeline for txt_seq_lens from version 0.37.0 to 0.39.0 across multiple Qwen image-related models Adds a new unit test to verify the deprecation warning behavior for the txt_seq_lens parameter * fix compile * formatting * fix compile tests * rename helper * remove duplicate * smaller values * removed * use torch.cond for torch compile * Construct joint attention mask once * test different backends * construct joint attention mask once to avoid reconstructing in every block * Update src/diffusers/models/attention_dispatch.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * formatting * raising an error from the EditPlus pipeline when batch_size > 1 --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com> Co-authored-by: YiYi Xu <yixu310@gmail.com> Co-authored-by: cdutr <dutra_carlos@hotmail.com>
1 parent b86bd99 commit dad5cb5

17 files changed

Lines changed: 513 additions & 172 deletions

‎docs/source/en/api/pipelines/qwenimage.md‎

Lines changed: 36 additions & 2 deletions

‎examples/dreambooth/train_dreambooth_lora_qwen_image.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1513,14 +1513,12 @@ def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
15131513
height=model_input.shape[3],
15141514
width=model_input.shape[4],
15151515
)
1516-
print(f"{prompt_embeds_mask.sum(dim=1).tolist()=}")
15171516
model_pred = transformer(
15181517
hidden_states=packed_noisy_model_input,
15191518
encoder_hidden_states=prompt_embeds,
15201519
encoder_hidden_states_mask=prompt_embeds_mask,
15211520
timestep=timesteps / 1000,
15221521
img_shapes=img_shapes,
1523-
txt_seq_lens=prompt_embeds_mask.sum(dim=1).tolist(),
15241522
return_dict=False,
15251523
)[0]
15261524
model_pred = QwenImagePipeline._unpack_latents(

‎src/diffusers/models/attention_dispatch.py‎

Lines changed: 76 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1945,6 +1945,43 @@ def score_mod(score, batch_idx, head_idx, q_idx, kv_idx):
19451945
return out
19461946

19471947

1948+
def _prepare_additive_attn_mask(
1949+
attn_mask: torch.Tensor, target_dtype: torch.dtype, reshape_4d: bool = True
1950+
) -> torch.Tensor:
1951+
"""
1952+
Convert a 2D attention mask to an additive mask, optionally reshaping to 4D for SDPA.
1953+
1954+
This helper is used by both native SDPA and xformers backends to handle both boolean and additive masks.
1955+
1956+
Args:
1957+
attn_mask: 2D tensor [batch_size, seq_len_k]
1958+
- Boolean: True means attend, False means mask out
1959+
- Additive: 0.0 means attend, -inf means mask out
1960+
target_dtype: The dtype to convert the mask to (usually query.dtype)
1961+
reshape_4d: If True, reshape from [batch_size, seq_len_k] to [batch_size, 1, 1, seq_len_k] for broadcasting
1962+
1963+
Returns:
1964+
Additive mask tensor where 0.0 means attend and -inf means mask out. Shape is [batch_size, seq_len_k] if
1965+
reshape_4d=False, or [batch_size, 1, 1, seq_len_k] if reshape_4d=True.
1966+
"""
1967+
# Check if the mask is boolean or already additive
1968+
if attn_mask.dtype == torch.bool:
1969+
# Convert boolean to additive: True -> 0.0, False -> -inf
1970+
attn_mask = torch.where(attn_mask, 0.0, float("-inf"))
1971+
# Convert to target dtype
1972+
attn_mask = attn_mask.to(dtype=target_dtype)
1973+
else:
1974+
# Already additive mask - just ensure correct dtype
1975+
attn_mask = attn_mask.to(dtype=target_dtype)
1976+
1977+
# Optionally reshape to 4D for broadcasting in attention mechanisms
1978+
if reshape_4d:
1979+
batch_size, seq_len_k = attn_mask.shape
1980+
attn_mask = attn_mask.view(batch_size, 1, 1, seq_len_k)
1981+
1982+
return attn_mask
1983+
1984+
19481985
@_AttentionBackendRegistry.register(
19491986
AttentionBackendName.NATIVE,
19501987
constraints=[_check_device, _check_shape],
@@ -1964,6 +2001,19 @@ def _native_attention(
19642001
) -> torch.Tensor:
19652002
if return_lse:
19662003
raise ValueError("Native attention backend does not support setting `return_lse=True`.")
2004+
2005+
# Reshape 2D mask to 4D for SDPA
2006+
# SDPA accepts both boolean masks (torch.bool) and additive masks (float)
2007+
if (
2008+
attn_mask is not None
2009+
and attn_mask.ndim == 2
2010+
and attn_mask.shape[0] == query.shape[0]
2011+
and attn_mask.shape[1] == key.shape[1]
2012+
):
2013+
# Just reshape [batch_size, seq_len_k] -> [batch_size, 1, 1, seq_len_k]
2014+
# SDPA handles both boolean and additive masks correctly
2015+
attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)
2016+
19672017
if _parallel_config is None:
19682018
query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
19692019
out = torch.nn.functional.scaled_dot_product_attention(
@@ -2530,10 +2580,34 @@ def _xformers_attention(
25302580
attn_mask = xops.LowerTriangularMask()
25312581
elif attn_mask is not None:
25322582
if attn_mask.ndim == 2:
2533-
attn_mask = attn_mask.view(attn_mask.size(0), 1, attn_mask.size(1), 1)
2583+
# Convert 2D mask to 4D for xformers
2584+
# Mask can be boolean (True=attend, False=mask) or additive (0.0=attend, -inf=mask)
2585+
# xformers requires 4D additive masks [batch, heads, seq_q, seq_k]
2586+
# Need memory alignment - create larger tensor and slice for alignment
2587+
original_seq_len = attn_mask.size(1)
2588+
aligned_seq_len = ((original_seq_len + 7) // 8) * 8 # Round up to multiple of 8
2589+
2590+
# Create aligned 4D tensor and slice to ensure proper memory layout
2591+
aligned_mask = torch.zeros(
2592+
(batch_size, num_heads_q, seq_len_q, aligned_seq_len),
2593+
dtype=query.dtype,
2594+
device=query.device,
2595+
)
2596+
# Convert to 4D additive mask (handles both boolean and additive inputs)
2597+
mask_additive = _prepare_additive_attn_mask(
2598+
attn_mask, target_dtype=query.dtype
2599+
) # [batch, 1, 1, seq_len_k]
2600+
# Broadcast to [batch, heads, seq_q, seq_len_k]
2601+
aligned_mask[:, :, :, :original_seq_len] = mask_additive
2602+
# Mask out the padding (already -inf from zeros -> where with default)
2603+
aligned_mask[:, :, :, original_seq_len:] = float("-inf")
2604+
2605+
# Slice to actual size with proper alignment
2606+
attn_mask = aligned_mask[:, :, :, :seq_len_kv]
25342607
elif attn_mask.ndim != 4:
25352608
raise ValueError("Only 2D and 4D attention masks are supported for xformers attention.")
2536-
attn_mask = attn_mask.expand(batch_size, num_heads_q, seq_len_q, seq_len_kv).type_as(query)
2609+
elif attn_mask.ndim == 4:
2610+
attn_mask = attn_mask.expand(batch_size, num_heads_q, seq_len_q, seq_len_kv).type_as(query)
25372611

25382612
if enable_gqa:
25392613
if num_heads_q % num_heads_kv != 0:

‎src/diffusers/models/controlnets/controlnet_qwenimage.py‎

Lines changed: 55 additions & 16 deletions

0 commit comments

Comments
 (0)