fix: device mismatch error in embedding loading on branch dev (#737)

* fix device mismatch bug

* fix additional device mismatch bug when zero-shot

* chore(format): run black on dev

---------

Co-authored-by: dragon <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: 3 people

ChatTTS/core.py

@@ -272,6 +272,7 @@ def _load(
                 vq_config=asdict(self.config.dvae.vq),
                 dim=self.config.dvae.decoder.idim,
                 coef=coef,
+                device=self.device,
             )
             .to(device)
             .eval()
@@ -288,8 +289,8 @@ def _load(
             self.config.embed.num_text_tokens,
             self.config.embed.num_vq,
         )
-        embed.from_pretrained(embed_path)
-        self.embed = embed
+        embed.from_pretrained(embed_path, device=self.device)
+        self.embed = embed.to(self.device)
         self.logger.log(logging.INFO, "embed loaded.")
 
         gpt = GPT(

ChatTTS/model/dvae.py

@@ -179,8 +179,10 @@ def __init__(
         hop_length=256,
         n_mels=100,
         padding: Literal["center", "same"] = "center",
+        device: torch.device = torch.device("cuda"),
     ):
         super().__init__()
+        self.device = device
         if padding not in ["center", "same"]:
             raise ValueError("Padding must be 'center' or 'same'.")
         self.padding = padding
@@ -197,6 +199,7 @@ 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)
         features = torch.log(torch.clip(mel, min=1e-5))
         return features
@@ -210,6 +213,7 @@ def __init__(
         vq_config: Optional[dict] = None,
         dim=512,
         coef: Optional[str] = None,
+        device: torch.device = torch.device("cuda"),
     ):
         super().__init__()
         if coef is None:
@@ -227,7 +231,7 @@ def __init__(
                 nn.Conv1d(dim, dim, 4, 2, 1),
                 nn.GELU(),
             )
-            self.preprocessor_mel = MelSpectrogramFeatures()
+            self.preprocessor_mel = MelSpectrogramFeatures(device=device)
             self.encoder: Optional[DVAEDecoder] = DVAEDecoder(**encoder_config)
 
         self.decoder = DVAEDecoder(**decoder_config)

ChatTTS/model/embed.py

@@ -34,12 +34,13 @@ def __init__(
         )
 
     @torch.inference_mode()
-    def from_pretrained(self, filename: str):
+    def from_pretrained(self, filename: str, device: torch.device):
         state_dict_tensors = {}
         with safe_open(filename, framework="pt") as f:
             for k in f.keys():
                 state_dict_tensors[k] = f.get_tensor(k)
         self.load_state_dict(state_dict_tensors)
+        self.to(device)
 
     def __call__(
         self, input_ids: torch.Tensor, text_mask: torch.Tensor