From 1e9e2be62281de76b9a761082d31c3654cd2b381 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 30 Nov 2025 17:32:24 +0200
Subject: [PATCH] Avoid recompile here

---
 wanvideo/modules/model.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py
index 87c2e21..1ae5100 100644
--- a/wanvideo/modules/model.py
+++ b/wanvideo/modules/model.py
@@ -412,11 +412,8 @@ class WanSelfAttention(nn.Module):
         v = self.v(x).view(b, s, n, d)
         return q, k, v
 
-    def qkv_fn_qk_with_rope(self, x, layer, freqs, num_chunks=1, is_longcat=False):
+    def _qkv_fn_with_rope(self, x, linear_layer, norm_layer, freqs, num_chunks=1, is_longcat=False):
         b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
-
-        linear_layer = self.q if layer == 'q' else self.k
-        norm_layer = self.norm_q if layer == 'q' else self.norm_k
         
         use_chunked = num_chunks > 1
         if use_chunked:
@@ -449,6 +446,12 @@ class WanSelfAttention(nn.Module):
                 result = norm_layer(linear_layer(x).to(norm_layer.weight.dtype)).to(x.dtype).view(b, s, n, d)
             return apply_rope_comfy1(result, freqs)
 
+    def qkv_fn_q_with_rope(self, x, freqs, num_chunks=1, is_longcat=False):
+        return self._qkv_fn_with_rope(x, self.q, self.norm_q, freqs, num_chunks, is_longcat)
+
+    def qkv_fn_k_with_rope(self, x, freqs, num_chunks=1, is_longcat=False):
+        return self._qkv_fn_with_rope(x, self.k, self.norm_k, freqs, num_chunks, is_longcat)
+
     def qkv_fn_v(self, x):
         b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
         return self.v(x).view(b, s, n, d)
@@ -1080,12 +1083,12 @@ class WanAttentionBlock(nn.Module):
             x_main, x_ip_input = input_x[:, : -self.cond_size], input_x[:, -self.cond_size :]
             # Compute QKV for main content
             if self.rope_func == "comfy":
-                q = self.self_attn.qkv_fn_qk_with_rope(x_main, "q", freqs)
-                k = self.self_attn.qkv_fn_qk_with_rope(x_main, "k", freqs)
+                q = self.self_attn.qkv_fn_q_with_rope(x_main, freqs)
+                k = self.self_attn.qkv_fn_k_with_rope(x_main, freqs)
                 v = self.self_attn.qkv_fn_v(x_main)
             elif self.rope_func == "comfy_chunked":
-                q = self.self_attn.qkv_fn_qk_with_rope(x_main, "q", freqs, num_chunks=2)
-                k = self.self_attn.qkv_fn_qk_with_rope(x_main, "k", freqs, num_chunks=2)
+                q = self.self_attn.qkv_fn_q_with_rope(x_main, freqs, num_chunks=2)
+                k = self.self_attn.qkv_fn_k_with_rope(x_main, freqs, num_chunks=2)
                 v = self.self_attn.qkv_fn_v(x_main)
             # Compute QKV for IP content
             if "comfy" in self.rope_func:
@@ -1094,8 +1097,8 @@ class WanAttentionBlock(nn.Module):
         else:
             if "comfy" in self.rope_func:
                 num_chunks = 2 if self.rope_func == "comfy_chunked" else 1
-                q = self.self_attn.qkv_fn_qk_with_rope(input_x, "q", freqs, num_chunks=num_chunks, is_longcat=is_longcat)
-                k = self.self_attn.qkv_fn_qk_with_rope(input_x, "k", freqs, num_chunks=num_chunks, is_longcat=is_longcat)
+                q = self.self_attn.qkv_fn_q_with_rope(input_x, freqs, num_chunks=num_chunks, is_longcat=is_longcat)
+                k = self.self_attn.qkv_fn_k_with_rope(input_x, freqs, num_chunks=num_chunks, is_longcat=is_longcat)
                 v = self.self_attn.qkv_fn_v(input_x)
             else:
                 q, k, v = self.self_attn.qkv_fn(input_x)