wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-02-09 04:52:32 +08:00
Code Issues Actions 9 Packages Projects Releases Wiki Activity

Fix Lumina model GQA compatibility with comfy_kitchen RoPE

Browse Source 
The comfy_kitchen apply_rope function requires xq and xk to have identical shapes, but Lumina uses Grouped Query Attention where query and key tensors can have different head counts. This adds a shape check to use apply_rope1 separately when shapes differ, maintaining compatibility with both standard attention and GQA.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
This commit is contained in:
keejkrej 2026-01-06 23:49:41 +00:00
parent 2c03884f5f
commit 98bdeca914
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
comfy/ldm/lumina/model.py
View File

@ -11,7 +11,7 @@ import comfy.ldm.common_dit
from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder
from comfy.ldm.modules.attention import optimized_attention_masked from comfy.ldm.modules.attention import optimized_attention_masked
from comfy.ldm.flux.layers import EmbedND from comfy.ldm.flux.layers import EmbedND
from comfy.ldm.flux.math import apply_rope from comfy.ldm.flux.math import apply_rope, apply_rope1
import comfy.patcher_extension import comfy.patcher_extension
@ -111,7 +111,12 @@ class JointAttention(nn.Module):
xq = self.q_norm(xq) xq = self.q_norm(xq)
xk = self.k_norm(xk) xk = self.k_norm(xk)
# Use apply_rope1 separately for GQA compatibility (when n_heads != n_kv_heads)
if xq.shape == xk.shape:
xq, xk = apply_rope(xq, xk, freqs_cis) xq, xk = apply_rope(xq, xk, freqs_cis)
else:
xq = apply_rope1(xq, freqs_cis)
xk = apply_rope1(xk, freqs_cis)
n_rep = self.n_local_heads // self.n_local_kv_heads n_rep = self.n_local_heads // self.n_local_kv_heads
if n_rep >= 1: if n_rep >= 1: