From 947c2749dd04c51ef0e21b069544d8b0b4f9b411 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Tue, 21 Jul 2026 20:02:45 -0700 Subject: [PATCH] Use optimized rms_rope function in joyai image model. (#15018) --- comfy/ldm/joyimage/model.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/comfy/ldm/joyimage/model.py b/comfy/ldm/joyimage/model.py index bca12c3917e..9d6951e54e4 100644 --- a/comfy/ldm/joyimage/model.py +++ b/comfy/ldm/joyimage/model.py @@ -94,12 +94,21 @@ def forward( txt_k = txt_k.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_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_k = torch.cat([img_k, txt_k], dim=1)