From b06b4f67cdd3e8e2bd9eff331623ab9c3674136f Mon Sep 17 00:00:00 2001 From: hben35096 <139383150+hben35096@users.noreply.github.com> Date: Fri, 9 May 2025 23:27:23 +0800 Subject: [PATCH] Solve the error "RuntimeError: BinaryCall MUDNN failed in: Run PowTensorOut" of the flux KSampler. --- comfy/ldm/flux/math.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/comfy/ldm/flux/math.py b/comfy/ldm/flux/math.py index 3e0978176..c1bfa32c8 100644 --- a/comfy/ldm/flux/math.py +++ b/comfy/ldm/flux/math.py @@ -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)