chore: restore some latest changes

2noise · Jul 21, 2024 · 9c0a1df · 9c0a1df
1 parent fe68af9
commit 9c0a1df
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/ChatTTS/core.py b/ChatTTS/core.py
@@ -253,6 +253,7 @@ def _load(
             device = select_device()
             self.logger.info("use device %s", str(device))
         self.device = device
+        self.device_gpt = device if "mps" not in str(device) else torch.device("cpu")
         self.compile = compile
 
         feature_extractor = instantiate_class(
@@ -299,7 +300,8 @@ def _load(
                 gpt = GPT(
                     **cfg,
                     use_flash_attn=use_flash_attn,
-                    device=device,
+                    device=self.device,
+                    device_gpt=self.device_gpt,
                     logger=self.logger,
                 ).eval()
                 assert gpt_ckpt_path, "gpt_ckpt_path should not be None"
@@ -537,7 +539,7 @@ def _infer_code(
             text,
             self.config.gpt.num_vq,
             prompt_str=params.spk_smp,
-            device=self.device,
+            device=self.device_gpt,
         )
         start_idx = input_ids.shape[-2]
 
@@ -597,7 +599,7 @@ def _refine_text(
         input_ids, attention_mask, text_mask = self.tokenizer.encode(
             text,
             self.config.gpt.num_vq,
-            device=self.device,
+            device=self.device_gpt,
         )
 
         start_idx = input_ids.shape[-2]

diff --git a/ChatTTS/model/gpt.py b/ChatTTS/model/gpt.py
@@ -28,14 +28,15 @@ def __init__(
         num_vq=4,
         use_flash_attn=False,
         device=torch.device("cpu"),
+        device_gpt=torch.device("cpu"),
         logger=logging.getLogger(__name__),
     ):
         super().__init__()
 
         self.logger = logger
 
         self.device = device
-        self.device_gpt = device if "mps" not in str(device) else torch.device("cpu")
+        self.device_gpt = device_gpt
 
         self.num_vq = num_vq
         self.num_audio_tokens = num_audio_tokens