[NPU][CosyVoice3] Add CPU STFT fallback (#7520)

Signed-off-by: dengchangxiong <dengchangxiong1@huawei.com>
Co-authored-by: dengchangxiong <dengchangxiong1@huawei.com>

tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py

@@ -76,6 +76,74 @@ def test_causal_hift_stft_moves_window_to_input_device(causal_hift):
     assert causal_hift.stft_window.device == waveform.device
 
 
+def test_causal_hift_routes_npu_transforms_to_cpu_fallback(causal_hift, monkeypatch):
+    class NpuTensor:
+        device = SimpleNamespace(type="npu")
+
+    tensor = NpuTensor()
+    stft_calls = []
+    istft_calls = []
+    monkeypatch.setattr(
+        causal_hift,
+        "_stft_on_cpu",
+        lambda value: stft_calls.append(value) or ("real", "imag"),
+    )
+    monkeypatch.setattr(
+        causal_hift,
+        "_istft_on_cpu",
+        lambda magnitude, phase: istft_calls.append((magnitude, phase)) or "waveform",
+    )
+
+    assert causal_hift._stft(tensor) == ("real", "imag")
+    assert causal_hift._istft(tensor, tensor) == "waveform"
+    assert stft_calls == [tensor]
+    assert istft_calls == [(tensor, tensor)]
+
+
+def test_causal_hift_stft_cpu_fallback_uses_float32(causal_hift, monkeypatch):
+    original_stft = torch.stft
+
+    def assert_cpu_float32_stft(input, *args, window, **kwargs):
+        assert input.device.type == "cpu"
+        assert input.dtype == torch.float32
+        assert window.device.type == "cpu"
+        assert window.dtype == torch.float32
+        return original_stft(input, *args, window=window, **kwargs)
+
+    monkeypatch.setattr(torch, "stft", assert_cpu_float32_stft)
+
+    waveform = torch.randn(1, 64, dtype=torch.float64)
+    real, imag = causal_hift._stft_on_cpu(waveform)
+
+    assert real.device == waveform.device
+    assert imag.device == waveform.device
+    assert real.dtype == waveform.dtype
+    assert imag.dtype == waveform.dtype
+
+
+def test_causal_hift_istft_cpu_fallback_uses_float32(causal_hift, monkeypatch):
+    original_istft = torch.istft
+
+    def assert_cpu_complex64_istft(input, *args, window, **kwargs):
+        assert input.device.type == "cpu"
+        assert input.dtype == torch.complex64
+        assert window.device.type == "cpu"
+        assert window.dtype == torch.float32
+        return original_istft(input, *args, window=window, **kwargs)
+
+    monkeypatch.setattr(torch, "istft", assert_cpu_complex64_istft)
+
+    waveform = torch.randn(1, 64, dtype=torch.float32)
+    spec = torch.stft(waveform, 16, 4, 16, window=causal_hift.stft_window, return_complex=True)
+    magnitude = torch.abs(spec).to(torch.float64)
+    phase = torch.angle(spec).to(torch.float64)
+    reconstructed = causal_hift._istft_on_cpu(magnitude, phase)
+
+    assert reconstructed.device == magnitude.device
+    assert reconstructed.dtype == magnitude.dtype
+    torch.testing.assert_close(reconstructed, waveform.to(torch.float64), rtol=1e-5, atol=1e-5)
+
+
 class TestPreLookaheadLayer:
     """Tests for PreLookaheadLayer."""
 

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py

@@ -489,6 +489,9 @@ class HiFTGenerator(nn.Module):
             block.remove_weight_norm()
 
     def _stft(self, x):
+        if x.device.type == "npu":
+            return self._stft_on_cpu(x)
+
         spec = torch.stft(
             x,
             self.istft_params["n_fft"],
@@ -500,18 +503,49 @@ class HiFTGenerator(nn.Module):
         spec = torch.view_as_real(spec)  # [B, F, TT, 2]
         return spec[..., 0], spec[..., 1]
 
+    def _stft_on_cpu(self, x):
+        target_device = x.device
+        target_dtype = x.dtype
+        spec = torch.stft(
+            x.float().cpu(),
+            self.istft_params["n_fft"],
+            self.istft_params["hop_len"],
+            self.istft_params["n_fft"],
+            window=self.stft_window.float().cpu(),
+            return_complex=True,
+        )
+        spec = torch.view_as_real(spec).to(device=target_device, dtype=target_dtype)
+        return spec[..., 0], spec[..., 1]
+
     def _istft(self, magnitude, phase):
+        if magnitude.device.type == "npu":
+            return self._istft_on_cpu(magnitude, phase)
+
         magnitude = torch.clip(magnitude, max=1e2)
         real = magnitude * torch.cos(phase)
-        img = magnitude * torch.sin(phase)
-        inverse_transform = torch.istft(
-            torch.complex(real, img),
+        imag = magnitude * torch.sin(phase)
+        return torch.istft(
+            torch.complex(real, imag),
             self.istft_params["n_fft"],
             self.istft_params["hop_len"],
             self.istft_params["n_fft"],
             window=self._get_stft_window(magnitude),
         )
-        return inverse_transform
+
+    def _istft_on_cpu(self, magnitude, phase):
+        target_device = magnitude.device
+        target_dtype = magnitude.dtype
+        magnitude = torch.clip(magnitude, max=1e2)
+        real = magnitude * torch.cos(phase)
+        imag = magnitude * torch.sin(phase)
+        waveform = torch.istft(
+            torch.complex(real.float().cpu(), imag.float().cpu()),
+            self.istft_params["n_fft"],
+            self.istft_params["hop_len"],
+            self.istft_params["n_fft"],
+            window=self.stft_window.float().cpu(),
+        )
+        return waveform.to(device=target_device, dtype=target_dtype)
 
     def _get_stft_window(self, tensor: torch.Tensor) -> torch.Tensor:
         if self.stft_window.device != tensor.device: