[BitcrushedHeart]

Summary

• Z-Image text caches pad all sequences to 512 tokens, but the padding is stripped at training time anyway (ZImageModel.encode_text() masks it out via tokens_mask)
• This adds two lightweight pipeline modules that strip padding before DiskCache saves and restore it after DiskCache loads, dramatically reducing text cache size on disk with no change to training behavior
• Only Z-Image is affected - no other model data loaders, mgds, or training code is modified

Details

• StripPaddedTokens (new, in modules/dataLoader/zimage/): inserted after EncodeQwenText in _preparation_modules(). Uses the attention mask to boolean-index tokens, tokens_mask, and text_encoder_hidden_state, keeping only real token positions before caching.
• PadTokens (new, in modules/dataLoader/zimage/): inserted at the start of _output_modules(). Pads stripped tensors back to PROMPT_MAX_LENGTH (512) using F.pad so torch.stack() batching still works. Old caches are compatible.

[BitcrushedHeart]

Tested. Works on Z-Image. Works with #1346 and Nerogar/mgds#44.

[dxqb]

Thanks for the PR!
Could you please have a look at this and check whether that changes anything for this PR? Tongyi-MAI/Z-Image#131 (comment)

Z-Image is unusual in that there are two types of padding tokens, inner padding and outer padding. Inner padding is not masked for attention (and must therefore not be cut!), outer padding is masked.

[BitcrushedHeart]

Thanks for the PR! Could you please have a look at this and check whether that changes anything for this PR? Tongyi-MAI/Z-Image#131 (comment)

Z-Image is unusual in that there are two types of padding tokens, inner padding and outer padding. Inner padding is not masked for attention (and must therefore not be cut!), outer padding is masked.

Thanks for flagging! I checked the linked issue & inner padding (which aligns sequences to SEQ_MULTI_OF) has attention_mask=1. StripPaddedTokens only removes positions where tokens_mask == 0 (outer padding), so inner padding is preserved. This matches the same boolean-mask filtering that encode_text() already applies at lines 171-172 of ZImageModel.py. No changes needed.

[dxqb]

Thanks.
Please see my code comments.
I think the same could work for other models, not only Z-Image:

• Chroma
• Qwen
• probably others

it won't work for

• SDXL
• Flux1
• Flux2

Would you like to implement it for those models as well?

[dxqb]

• update mgds Add PruneMaskedTokens and PadMaskedTokens pipeline modules mgds#47

[dxqb]

looks good to me, but I am a bit hesitant to merge this: any silent errors in this PR will be difficult to detect but screws up text conditioning and have very unexpected effects.

any way this can be thoroughly tested?

[dxqb]

for a small dataset captioned by an AI, this reduced total cache size from 98 MB to 14 MB on Qwen

[dxqb]

tested for all 3 models

identical for z-Image, identical for unmasked tokens on the other models (but if there isn't a bug in diffusers for attention masking, that's okay)

can be merged once Nerogar/mgds#47 is finished

here is how I tested:

                filename = f"text-{train_progress.global_step}.pt"
                if False:
                    torch.save((text_encoder_output, text_attention_mask), filename)
                else:
                    text_encoder_output2, text_attention_mask2 = torch.load(filename)
                    identical = torch.equal(text_encoder_output, text_encoder_output2)
                    masked = torch.equal(text_encoder_output[text_attention_mask], text_encoder_output2[text_attention_mask2])
                    print("\n", identical, masked)
                    assert masked