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

Use optimized rms_rope function in joyai image model. (#15018)

Browse Source
This commit is contained in:
comfyanonymous
2026-07-21 20:02:45 -07:00
committed by GitHub
parent 7bf8bfcd07
commit 947c2749dd
1 changed files with 12 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
comfy/ldm/joyimage/model.py
View File

@@ -94,12 +94,21 @@ class JoyImageAttention(nn.Module):
txt_k = txt_k.unflatten(-1, (heads, -1)) txt_k = txt_k.unflatten(-1, (heads, -1))
txt_v = txt_v.unflatten(-1, (heads, -1)) txt_v = txt_v.unflatten(-1, (heads, -1))
img_q = self.img_attn_q_norm(img_q)
img_k = self.img_attn_k_norm(img_k)
txt_q = self.txt_attn_q_norm(txt_q) txt_q = self.txt_attn_q_norm(txt_q)
txt_k = self.txt_attn_k_norm(txt_k) txt_k = self.txt_attn_k_norm(txt_k)
img_q, img_k = comfy_kitchen.apply_rope(img_q, img_k, image_rotary_emb) img_q_scale, _, img_q_offload_stream = comfy.ops.cast_bias_weight(self.img_attn_q_norm, img_q, offloadable=True)
img_k_scale, _, img_k_offload_stream = comfy.ops.cast_bias_weight(self.img_attn_k_norm, img_k, offloadable=True)
img_q, img_k = comfy_kitchen.rms_rope(
img_q,
img_k,
image_rotary_emb,
img_q_scale,
img_k_scale,
self.img_attn_q_norm.eps,
)
comfy.ops.uncast_bias_weight(self.img_attn_q_norm, img_q_scale, None, img_q_offload_stream)
comfy.ops.uncast_bias_weight(self.img_attn_k_norm, img_k_scale, None, img_k_offload_stream)
joint_q = torch.cat([img_q, txt_q], dim=1) joint_q = torch.cat([img_q, txt_q], dim=1)
joint_k = torch.cat([img_k, txt_k], dim=1) joint_k = torch.cat([img_k, txt_k], dim=1)