Z-Image text sequence length issue

[dxqb]

Describe the bug

I think there might be an issue with calculating the sequence length of cap_feat (which is the text encoder output), and masking it accordingly.

I'm going to use code links from before the omni commit, because it's easier to read - but the issue seems exist in both, before and after the omni commit.

diffusers/src/diffusers/models/transformers/transformer_z_image.py

Lines 593 to 611 in 52766e6

	cap_item_seqlens = [len(_) for _ in cap_feats]
	cap_max_item_seqlen = max(cap_item_seqlens)
	cap_feats = torch.cat(cap_feats, dim=0)
	cap_feats = self.cap_embedder(cap_feats)
	cap_feats[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token
	cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0))
	cap_freqs_cis = list(
	self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split([len(_) for _ in cap_pos_ids], dim=0)
	)
	cap_feats = pad_sequence(cap_feats, batch_first=True, padding_value=0.0)
	cap_freqs_cis = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0)
	# Clarify the length matches to satisfy Dynamo due to "Symbolic Shape Inference" to avoid compilation errors
	cap_freqs_cis = cap_freqs_cis[:, : cap_feats.shape[1]]
	cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device)
	for i, seq_len in enumerate(cap_item_seqlens):
	cap_attn_mask[i, :seq_len] = 1

In these lines, the sequence length for each cap_feats sample is put into cap_item_seqlens and is then used to create an attention mask in line 611.

The problem with this is that cap_feats has been overwritten before in these lines:

diffusers/src/diffusers/models/transformers/transformer_z_image.py

Lines 552 to 560 in 52766e6

	(
	x,
	cap_feats,
	x_size,
	x_pos_ids,
	cap_pos_ids,
	x_inner_pad_mask,
	cap_inner_pad_mask,
	) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size)

cap_feats therefore isn't anymore what the caller has passed. instead, cap_feats has already been padded by patchify_and_embed. Therefore, all cap_item_seqlens are identical to cap_max_item_seqlen.

This leads to unmasked text tokens during attention, which probably wasn't the intention here.

Reproduction

put a breakpoint in this line:

diffusers/src/diffusers/models/transformers/transformer_z_image.py

Line 612 in 52766e6

all cap_items_seqlens are identical, cap_attn_mask is all True - even if multiple different text lengths were passed to the transformer.

Logs

System Info

Who can help?

@JerryWu-code @RuoyiDu @yiyixuxu

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[dxqb]

this should still be adressed @JerryWu-code @RuoyiDu

[Mr-Neutr0n]

Opened a fix for this in #13131.

The root cause is that _prepare_sequence computes item_seqlens from the already-padded features (padded to multiples of SEQ_MULTI_OF by _pad_with_ids), so the attention mask treats all inner padding positions as valid. The fix uses the existing inner_pad_mask to exclude those positions from the attention mask, both in the individual refiner stages and in the unified sequence passed to the main transformer layers.

[ChrisLiu6]

Hi guys, could you kindly take a look at this issue Tongyi-MAI/Z-Image#131 and check if this solves your questions?

[dxqb]

please ignore my previous (now deleted) message. this was before I understood your answer.

Hi guys, could you kindly take a look at this issue Tongyi-MAI/Z-Image#131 and check if this solves your questions?

Thank you. That's a smart idea.
I guess to avoid masking, but still be able to have variable caption lengths in the same batch during training. And without having a mismatch between training and inference by training the model to deal with the "inner padding".

The only thing I'd suggest then is to actually use the benefit of this. The current code always uses the attention mask in attention, even if it's all-True and therefore a no-op. This causes torch SDPA to not use its internal flash attention algorithm, but the slower one that supports masks.

I have prepared a small PR here: #13184
Please let me know if you agree.

It's about 15% faster for inference(even for batch size 1). Usually more in training but I haven't benchmarked it now.
The inference result is pixel-identical.

[ChrisLiu6]

Thanks! The PR looks good to me.

[Beinsezii]

Probably also fixed by #12955

[yiyixuxu]

should be fixed now we merged #12955