fix bugs of npu support

2noise · Oct 16, 2024 · 6d20ce0 · 6d20ce0
1 parent 78ecb3c
commit 6d20ce0
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/ChatTTS/model/dvae.py b/ChatTTS/model/dvae.py
@@ -199,8 +199,15 @@ def __call__(self, audio: torch.Tensor) -> torch.Tensor:
         return super().__call__(audio)
 
     def forward(self, audio: torch.Tensor) -> torch.Tensor:
-        audio = audio.to(self.device)
-        mel: torch.Tensor = self.mel_spec(audio)
+        if "npu" in str(self.device):
+            # Computation of MelSpectrogram on npu is not supported now, use cpu fallback.
+            audio = audio.to(torch.device("cpu"))
+            self.mel_spec.to(torch.device("cpu"))
+            mel: torch.Tensor = self.mel_spec(audio)
+            mel = mel.to(self.device)
+        else:
+            audio = audio.to(self.device)
+            mel: torch.Tensor = self.mel_spec(audio)
         features = torch.log(torch.clip(mel, min=1e-5))
         return features