Repository navigation
Z-Image的文本位置编码长度不一致问题 #13574
Description
# Process captions
for j, cap_item in enumerate(all_cap_feats[i]):
noise_val = images_noise_mask[i][j] if j < len(images_noise_mask[i]) else 1
cap_out, cap_pos, cap_mask, cap_len, cap_nm = self._pad_with_ids(
cap_item,
(len(cap_item) + (-len(cap_item)) % SEQ_MULTI_OF, 1, 1),
(cap_cu_len, 0, 0),
device,
noise_val,
)
cap_feats_list.append(cap_out)
cap_pos_list.append(cap_pos)
cap_mask_list.append(cap_mask)
cap_lens.append(cap_len)
cap_noise.extend(cap_nm)
cap_cu_len += len(cap_item)
cap_end_pos.append(cap_cu_len)
cap_cu_len += 2 # for image vae and siglip tokens
这里cap_out和cap_pos输出的长度不一样？假设len(cap_item)=120，cap_out输出是128，cap_pos是128+8=136
Activity
chinmayirm commentedon Sep 22, 2026
Hi @PlutoQyl, the length mismatch you pasted is real.
Root cause
_pad_with_idsbuilds a coordinate grid ofpos_grid_size, then, when the feature is short of a multiple of 32, appendspad_lenmore position ids. The caption call sites already pass the padded length as that grid:transformer_z_image.py, in bothpatchify_and_embedandpatchify_and_embed_omni(the same call is copied incontrolnet_z_image.py):(len(cap_item) + (-len(cap_item)) % SEQ_MULTI_OF, 1, 1)For
len(cap_item) == 120,pad_lenis 8. The grid is 128 positions, the function then appends 8 more, andcap_posis 136 whilecap_outis 128. Image tokens do not hit this: their grid is(F_tokens, H_tokens, W_tokens), which matches the unpadded patch count, and the pad is added once._prepare_sequencethen doesfreqs_cis[:, : feats.shape[1]], so the extra 8 ids are dropped and it does not crash. The ids that survive are the first 128 of that padded grid. Real caption tokens keep the right coordinates. The pad tokens keep the next coordinates instead of the(0, 0, 0)ids_pad_with_idsjust built, and those are the ones the slice throws away.Proposed approach
Pass the unpadded caption length into the grid,
(len(cap_item), 1, 1), at all three call sites._pad_with_idsalready appends the pad ids.cap_outandcap_posthen both have length 128, the slice is a no-op, and pad tokens sit on(0, 0, 0), which is what that function writes. Real-token coordinates stay the same, because those are the firstori_lenentries of either grid.Tests
On a caption length that is not a multiple of 32:
cap_out.shape[0] == cap_pos.shape[0] == cap_pad_mask.shape[0], the firstori_lenposition ids match a grid of the unpadded length, and the pad rows are(0, 0, 0). One case that is already a multiple of 32, wherepad_lenis 0. The same assertions on the copied ControlNet path.I would be glad to open a PR for this once the approach looks right to you.