upstream-archive/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-07-23 09:10:51 -05:00
Code Issues Packages Projects Releases Wiki Activity

Make z image/lumina 2 models use comfy kitchen rms rope. (#15036)

Browse Source
This commit is contained in:
comfyanonymous
2026-07-22 12:34:27 -07:00
committed by GitHub
parent 54ca9193a3
commit 2e47082c8e
1 changed files with 19 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -6,6 +6,9 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
import comfy.ldm.common_dit
import comfy.model_management
import comfy.ops
import comfy.quant_ops
from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder
from comfy.ldm.modules.attention import optimized_attention_masked
@@ -97,6 +100,7 @@ class JointAttention(nn.Module):
self.n_local_kv_heads = self.n_kv_heads
self.n_rep = self.n_local_heads // self.n_local_kv_heads
self.head_dim = dim // n_heads
self.qk_norm = qk_norm
self.qkv = operation_settings.get("operations").Linear(
dim,
@@ -151,10 +155,21 @@ class JointAttention(nn.Module):
xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
xq = self.q_norm(xq)
xk = self.k_norm(xk)
xq, xk = apply_rope(xq, xk, freqs_cis)
if self.qk_norm and not comfy.model_management.in_training:
q_scale, _, q_offload_stream = comfy.ops.cast_bias_weight(self.q_norm, xq, offloadable=True)
k_scale, _, k_offload_stream = comfy.ops.cast_bias_weight(self.k_norm, xk, offloadable=True)
epsilon = self.q_norm.eps if self.q_norm.eps is not None else torch.finfo(torch.float32).eps
if self.n_local_heads == self.n_local_kv_heads:
xq, xk = comfy.quant_ops.ck.rms_rope(xq, xk, freqs_cis, q_scale, k_scale, epsilon)
else:
xq = comfy.quant_ops.ck.rms_rope1(xq, freqs_cis, q_scale, epsilon)
xk = comfy.quant_ops.ck.rms_rope1(xk, freqs_cis, k_scale, epsilon)
comfy.ops.uncast_bias_weight(self.q_norm, q_scale, None, q_offload_stream)
comfy.ops.uncast_bias_weight(self.k_norm, k_scale, None, k_offload_stream)
else:
xq = self.q_norm(xq)
xk = self.k_norm(xk)
xq, xk = apply_rope(xq, xk, freqs_cis)
n_rep = self.n_local_heads // self.n_local_kv_heads
if n_rep >= 1: