Fix rope encoding.

This commit is contained in:
ZHOU He
2025-11-16 14:48:45 +08:00
parent f5163aab75
commit ebd760c92f
+2 -5
View File
@@ -886,11 +886,8 @@ class GenConsistent:
if pe is not None:
# Extract positional encoding for the saved key positions
# This mimics the approach in attention_processor_characonsist.py
image_rotary_emb_fg = (
pe[0][valid_curr_inds] if pe[0].shape[0] >= len(valid_curr_inds) else pe[0],
pe[1][valid_curr_inds] if pe[1].shape[0] >= len(valid_curr_inds) else pe[1],
)
saved_key_fg = flux_math.apply_rope(saved_key_fg, image_rotary_emb_fg)
image_rotary_emb_fg = pe[:,:,valid_curr_inds]
saved_key_fg = flux_math.apply_rope1(saved_key_fg, image_rotary_emb_fg)
# Concatenate the expanded key-value pairs
k = torch.cat([k, saved_key_fg], dim=2)