RomanNum3ral
/
ComfyUI
mirror of https://github.com/Comfy-Org/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Make Ernie use comfy kitchen rms rope (#15055)

This commit is contained in:
comfyanonymous 2026-07-23 19:06:52 -07:00 committed by GitHub
parent feca51a854
commit 0cb84e7e6e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 12 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
comfy/ldm/ernie/model.py
View File

@ -5,6 +5,7 @@ import torch.nn.functional as F
from comfy.ldm.modules.attention import optimized_attention
import comfy.model_management
import comfy.ops
import comfy.quant_ops
def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
@ -111,11 +112,17 @@ class ErnieImageAttention(nn.Module):
query = q_flat.view(B, S, self.heads, self.head_dim)
key = k_flat.view(B, S, self.heads, self.head_dim)
query = self.norm_q(query)
key = self.norm_k(key)
if image_rotary_emb is not None:
query, key = comfy.quant_ops.ck.apply_rope_split_half(query, key, image_rotary_emb)
if image_rotary_emb is not None and not comfy.model_management.in_training:
q_scale, _, q_offload_stream = comfy.ops.cast_bias_weight(self.norm_q, query, offloadable=True)
k_scale, _, k_offload_stream = comfy.ops.cast_bias_weight(self.norm_k, key, offloadable=True)
query, key = comfy.quant_ops.ck.rms_rope_split_half(query, key, image_rotary_emb, q_scale, k_scale, self.norm_q.eps)
comfy.ops.uncast_bias_weight(self.norm_q, q_scale, None, q_offload_stream)
comfy.ops.uncast_bias_weight(self.norm_k, k_scale, None, k_offload_stream)
else:
query = self.norm_q(query)
key = self.norm_k(key)
if image_rotary_emb is not None:
query, key = comfy.quant_ops.ck.apply_rope_split_half(query, key, image_rotary_emb)
q_flat = query.reshape(B, S, -1)
k_flat = key.reshape(B, S, -1)