perf: use cpu in MacOS by default

2noise · Jun 27, 2024 · 070a0d3 · 070a0d3
1 parent 0716d56
commit 070a0d3
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 11 deletions.
diff --git a/ChatTTS/core.py b/ChatTTS/core.py
@@ -438,6 +438,14 @@ def _text_to_token(self, text: str, device="cpu") -> Tuple[torch.Tensor, torch.T
         del_all(text_token)
 
         return input_ids, attention_mask, text_mask
+
+    @staticmethod
+    def _decode_spk_emb(spk_emb: str) -> np.ndarray:
+        return np.frombuffer(lzma.decompress(
+            b14.decode_from_string(spk_emb),
+            format=lzma.FORMAT_RAW,
+            filters=[{"id": lzma.FILTER_LZMA2, "preset": 9 | lzma.PRESET_EXTREME}],
+        ), dtype=np.float16).copy()
 
     def _apply_spk_emb(
         self,
@@ -448,11 +456,7 @@ def _apply_spk_emb(
     ):
         n = F.normalize(
             torch.from_numpy(
-                np.frombuffer(lzma.decompress(
-                    b14.decode_from_string(spk_emb),
-                    format=lzma.FORMAT_RAW,
-                    filters=[{"id": lzma.FILTER_LZMA2, "preset": 9 | lzma.PRESET_EXTREME}],
-                ), dtype=np.float16).copy(),
+                self._decode_spk_emb(spk_emb),
             ).unsqueeze(0).expand(text_len, -1), p=2.0, dim=1, eps=1e-12
         ).to(self.gpt.device_gpt).unsqueeze_(1).expand(emb.shape)
         cond = input_ids.narrow(-1, 0, 1).eq(self.tokenizer_spk_emb_ids).expand(emb.shape)

diff --git a/ChatTTS/utils/gpu.py b/ChatTTS/utils/gpu.py
@@ -3,7 +3,7 @@
 from .log import logger
 
 
-def select_device(min_memory=2047):
+def select_device(min_memory=2047, experimental=False):
     if torch.cuda.is_available():
         available_gpus = []
         for i in range(torch.cuda.device_count()):
@@ -19,9 +19,17 @@ def select_device(min_memory=2047):
             )
             device = torch.device("cpu")
     elif torch.backends.mps.is_available():
-        # For Apple M1/M2 chips with Metal Performance Shaders
-        logger.get_logger().info("apple GPU found, using MPS.")
-        device = torch.device("mps")
+        """
+        Currently MPS is slower than CPU while needs more memory and core utility,
+        so only enable this for experimental use.
+        """
+        if experimental:
+            # For Apple M1/M2 chips with Metal Performance Shaders
+            logger.get_logger().warn("experimantal: found apple GPU, using MPS.")
+            device = torch.device("mps")
+        else:
+            logger.get_logger().info("found Apple GPU, but use CPU.")
+            device = torch.device("cpu")
     else:
         logger.get_logger().warning("no GPU found, use CPU instead")
         device = torch.device("cpu")

diff --git a/examples/cmd/run.py b/examples/cmd/run.py
@@ -6,9 +6,7 @@
 now_dir = os.getcwd()
 sys.path.append(now_dir)
 
-import wave
 import argparse
-from io import BytesIO
 
 import ChatTTS