CodeCow/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-07-21 21:01:24 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

rope

Browse Source
This commit is contained in:
Haoming 2026-02-19 00:46:30 +08:00
parent 62828826b3
commit c3629f82d2
1 changed files with 12 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
backend/nn/flux.py
View File

@ -37,11 +37,18 @@ def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
return out.to(dtype=torch.float32, device=pos.device)
def apply_rope1(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
x_ = x.to(dtype=freqs_cis.dtype).reshape(*x.shape[:-1], -1, 1, 2)
x_out = freqs_cis[..., 0] * x_[..., 0]
x_out.addcmul_(freqs_cis[..., 1], x_[..., 1])
return x_out.reshape(*x.shape).type_as(x)
try:
import comfy_kitchen as ck
except ImportError:
def apply_rope1(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
x_ = x.to(dtype=freqs_cis.dtype).reshape(*x.shape[:-1], -1, 1, 2)
x_out = freqs_cis[..., 0] * x_[..., 0]
x_out.addcmul_(freqs_cis[..., 1], x_[..., 1])
return x_out.reshape(*x.shape).type_as(x)
else:
apply_rope1 = ck.apply_rope1
def apply_rope(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: