From f208f907bc351dbb863e353ab5d996707bd4c339 Mon Sep 17 00:00:00 2001 From: Weiming Liao Date: Wed, 19 Aug 2026 00:06:29 +0800 Subject: [PATCH] [BugFix][NPU] Wan2.2-S2V RoPE: replace complex64 advanced indexing with index_select (#6320) Signed-off-by: Weiming Liao --- vllm_omni/diffusion/layers/rope.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index fae1a3414..e78574b04 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -423,16 +423,16 @@ class WanS2VRotaryPosEmbed(torch.nn.Module): h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long() w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long() - freqs_0 = freqs_split[0][f_sam] if f_o >= 0 else freqs_split[0][f_sam].conj() + freqs_0 = torch.index_select(freqs_split[0] if f_o >= 0 else freqs_split[0].conj(), 0, f_sam) freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1) freqs_i = torch.cat( [ freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1), - freqs_split[1][h_sam] + torch.index_select(freqs_split[1], 0, h_sam) .view(1, seq_h_int, 1, -1) .expand(seq_f_int, seq_h_int, seq_w_int, -1), - freqs_split[2][w_sam] + torch.index_select(freqs_split[2], 0, w_sam) .view(1, 1, seq_w_int, -1) .expand(seq_f_int, seq_h_int, seq_w_int, -1), ], @@ -523,14 +523,18 @@ class RotaryEmbeddingS2VGrid(torch.nn.Module): h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long() w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long() - freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][f_sam].conj() + freqs_0 = torch.index_select(freqs[0] if f_o >= 0 else freqs[0].conj(), 0, f_sam) freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1) freqs_i = torch.cat( [ freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1), - freqs[1][h_sam].view(1, seq_h_int, 1, -1).expand(seq_f_int, seq_h_int, seq_w_int, -1), - freqs[2][w_sam].view(1, 1, seq_w_int, -1).expand(seq_f_int, seq_h_int, seq_w_int, -1), + torch.index_select(freqs[1], 0, h_sam) + .view(1, seq_h_int, 1, -1) + .expand(seq_f_int, seq_h_int, seq_w_int, -1), + torch.index_select(freqs[2], 0, w_sam) + .view(1, 1, seq_w_int, -1) + .expand(seq_f_int, seq_h_int, seq_w_int, -1), ], dim=-1, ).reshape(seg_len, 1, -1)