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

comfyanonymous · web-flow · commit 947c2749dd04 · 2026-07-21T20:02:45.000-07:00
diff --git 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)
