xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-02 21:07:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

Solve the error "RuntimeError: BinaryCall MUDNN failed in: Run PowTensorOut" of the flux KSampler.

Browse Source
This commit is contained in:
hben35096 2025-05-09 23:27:23 +08:00 committed by GitHub
parent d22feacb05
commit b06b4f67cd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
comfy/ldm/flux/math.py
View File

@ -27,9 +27,14 @@ def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
device = torch.device("cpu")
else:
device = pos.device
scale = torch.linspace(0, (dim - 2) / dim, steps=dim//2, dtype=torch.float64, device=device)
omega = 1.0 / (theta**scale)
if device.type == "musa":
scale = torch.linspace(0, (dim - 2) / dim, steps=dim // 2, dtype=torch.float32, device=device)
if not isinstance(theta, torch.Tensor):
theta = torch.tensor(theta, dtype=torch.float32, device=device)
omega = torch.exp(-scale * torch.log(theta + 1e-6))
else:
scale = torch.linspace(0, (dim - 2) / dim, steps=dim//2, dtype=torch.float64, device=device)
omega = 1.0 / (theta**scale)
out = torch.einsum("...n,d->...nd", pos.to(dtype=torch.float32, device=device), omega)
out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)