From ebd760c92fc3cd5d4d521335c1e153ed42766e94 Mon Sep 17 00:00:00 2001
From: ZHOU He <461607353@qq.com>
Date: Sun, 16 Nov 2025 14:48:45 +0800
Subject: [PATCH] Fix rope encoding.

---
 node.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/node.py b/node.py
index 0110f2d..296a521 100644
--- a/node.py
+++ b/node.py
@@ -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)