wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-06-15 04:19:43 +08:00
Code Issues Actions 18 Packages Projects Releases Wiki Activity

Use comfy kitchen apply rope in omnigen2 model. (#14442)

Browse Source
This commit is contained in:
comfyanonymous 2026-06-12 18:38:39 -07:00 committed by GitHub
parent 28a40fb2b2
commit 7277d99d3a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
comfy/ldm/omnigen/omnigen2.py
View File

@ -8,6 +8,7 @@ import torch.nn.functional as F
from einops import rearrange, repeat from einops import rearrange, repeat
from comfy.ldm.lightricks.model import Timesteps from comfy.ldm.lightricks.model import Timesteps
from comfy.ldm.flux.layers import EmbedND from comfy.ldm.flux.layers import EmbedND
from comfy.ldm.flux.math import apply_rope1
from comfy.ldm.modules.attention import optimized_attention_masked from comfy.ldm.modules.attention import optimized_attention_masked
import comfy.model_management import comfy.model_management
import comfy.ldm.common_dit import comfy.ldm.common_dit
@ -17,9 +18,7 @@ def apply_rotary_emb(x, freqs_cis):
if x.shape[1] == 0: if x.shape[1] == 0:
return x return x
t_ = x.reshape(*x.shape[:-1], -1, 1, 2) return apply_rope1(x, freqs_cis)
t_out = freqs_cis[..., 0] * t_[..., 0] + freqs_cis[..., 1] * t_[..., 1]
return t_out.reshape(*x.shape).to(dtype=x.dtype)
def swiglu(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: def swiglu(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: