From b366a40ab6ec0e8529c6309c058e69e8773004d3 Mon Sep 17 00:00:00 2001 From: zhuhz22 <138758655+zhuhz22@users.noreply.github.com> Date: Sun, 2 Mar 2025 14:38:08 +0800 Subject: [PATCH] fix: riflex --- wanvideo/modules/model.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py index 79e062d..bcb86b7 100644 --- a/wanvideo/modules/model.py +++ b/wanvideo/modules/model.py @@ -36,19 +36,19 @@ def sinusoidal_embedding_1d(dim, position): return x -def rope_params(max_seq_len, dim, theta=10000, L_test=81, k=0): +def rope_params(max_seq_len, dim, theta=10000, L_test=25, k=0): assert dim % 2 == 0 - freqs = torch.outer( - torch.arange(max_seq_len), - 1.0 / torch.pow(theta, - torch.arange(0, dim, 2).to(torch.float64).div(dim))) + exponents = torch.arange(0, dim, 2, dtype=torch.float64).div(dim) + inv_theta_pow = 1.0 / torch.pow(theta, exponents) + if k > 0: print(f"RifleX: Using {k}th freq") - freqs[k-1] = 0.9 * 2 * torch.pi / L_test + inv_theta_pow[k-1] = 0.9 * 2 * torch.pi / L_test + + freqs = torch.outer(torch.arange(max_seq_len), inv_theta_pow) freqs = torch.polar(torch.ones_like(freqs), freqs) return freqs - from comfy.model_management import get_torch_device, get_autocast_device @torch.autocast(device_type=get_autocast_device(get_torch_device()), enabled=False) @torch.compiler.disable()