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)