Fix QwenImage txt_seq_lens handling #12702
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
just a few comments, not a full review:
|
| raise ValueError(f"`txt_seq_lens` must have length {batch_size}, but got {len(txt_seq_lens)} instead.") | ||
| text_seq_len = max(text_seq_len, max(txt_seq_lens)) | ||
| elif encoder_hidden_states_mask is not None: | ||
| text_seq_len = max(text_seq_len, int(encoder_hidden_states_mask.sum(dim=1).max().item())) |
There was a problem hiding this comment.
This only works if the attention mask is in the form of [True, True, True, ..., False, False, False]. While this is the case in the most common use case of text attention masks, it doesn't have to be the case.
If the mask is [True, False, True, False, True, False], self.pos_embed receives an incorrect sequence length
|
thanks @dxqb the idea was to remove |
- 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
|
some benchmarks with various backends: code: 
|
|
cc @naykun can you take a look here too? |
Co-authored-by: YiYi Xu <yixu310@gmail.com>
Co-authored-by: YiYi Xu <yixu310@gmail.com>
Enhances documentation with comprehensive performance insights for QwenImage pipeline:
|
Hey @kashif! I've prepared a documentation update with a new Performance section covering:
I am also mentioning the gist with the scripts I used Can you double check? Is there anything missing, or something else I can help with? |
| if self.scale_rope: | ||
| max_vid_index = max(height // 2, width // 2, max_vid_index) | ||
| else: | ||
| max_vid_index = max(height, width, max_vid_index) |
There was a problem hiding this comment.
If we want to support batch inference in this PR, further modifications are needed. The current implementation assumes all instances have the same size, but in general cases—especially in editing scenarios—input images may vary in size. To handle this, we should compute max_vid_index per instance and slice txt_freqs accordingly.
Otherwise, a warning should be raised to alert users who are running batch inference.
| # Keep as tensor to avoid graph breaks in torch.compile | ||
| # torch.maximum works with mixed tensor/scalar and keeps result as tensor | ||
| text_seq_len_tensor = torch.tensor(text_seq_len, device=encoder_hidden_states.device, dtype=torch.long) | ||
| rope_text_seq_len = torch.maximum(text_seq_len_tensor, per_sample_len.max()) |
There was a problem hiding this comment.
Will rope_text_seq_len always be equal to text_seq_len_tensor?
| device: torch.device, | ||
| txt_seq_lens: Optional[List[int]] = None, | ||
| device: torch.device = None, | ||
| txt_seq_len: Optional[Union[int, torch.Tensor]] = None, |
There was a problem hiding this comment.
This name could be confusing due to the existing txt_seq_lens parameter. Would max_txt_seq_len be clearer for txt_seq_len?
|
Thank you so much for this excellent PR! It’s clean, well-structured, and addresses several long-standing issues. I’ve left a few questions —we can discuss them further. |
7 participants
What does this PR do?
txt_seq_lensplumbing from all QwenImage pipelines and modular steps; the transformer now infers text length from encoder inputs/masks and validates optional overrides.encoder_hidden_states_maskinside the double-stream attention, avoiding fullseq_len²masks while keeping padding tokens masked.txt_seq_lensvalues and encoder masks are handled safely.Fixes #12344
Before submitting
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.