From 7277d99d3ab9f4ced8db4e87c82bad68028aa80d Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Fri, 12 Jun 2026 18:38:39 -0700 Subject: [PATCH] Use comfy kitchen apply rope in omnigen2 model. (#14442) --- comfy/ldm/omnigen/omnigen2.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/comfy/ldm/omnigen/omnigen2.py b/comfy/ldm/omnigen/omnigen2.py index 82edc92da..e9ca5229d 100644 --- a/comfy/ldm/omnigen/omnigen2.py +++ b/comfy/ldm/omnigen/omnigen2.py @@ -8,6 +8,7 @@ import torch.nn.functional as F from einops import rearrange, repeat from comfy.ldm.lightricks.model import Timesteps from comfy.ldm.flux.layers import EmbedND +from comfy.ldm.flux.math import apply_rope1 from comfy.ldm.modules.attention import optimized_attention_masked import comfy.model_management import comfy.ldm.common_dit @@ -17,9 +18,7 @@ def apply_rotary_emb(x, freqs_cis): if x.shape[1] == 0: return x - t_ = x.reshape(*x.shape[:-1], -1, 1, 2) - t_out = freqs_cis[..., 0] * t_[..., 0] + freqs_cis[..., 1] * t_[..., 1] - return t_out.reshape(*x.shape).to(dtype=x.dtype) + return apply_rope1(x, freqs_cis) def swiglu(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: