Fix coqui-ai#2749 (coqui-ai#2750)

pugtech-co · Aug 14, 2023 · 75f0d13 · 75f0d13
1 parent 4bfda13
commit 75f0d13
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/TTS/tts/layers/bark/hubert/tokenizer.py b/TTS/tts/layers/bark/hubert/tokenizer.py
@@ -121,7 +121,7 @@ def load_from_checkpoint(path, map_location: MAP_LOCATION = None):
                 data_from_model.output_size,
                 data_from_model.version,
             )
-        model.load_state_dict(torch.load(path))
+        model.load_state_dict(torch.load(path, map_location=map_location))
         if map_location:
             model = model.to(map_location)
         return model

diff --git a/TTS/tts/layers/bark/inference_funcs.py b/TTS/tts/layers/bark/inference_funcs.py
@@ -136,9 +136,7 @@ def generate_voice(
     hubert_model = CustomHubert(checkpoint_path=model.config.LOCAL_MODEL_PATHS["hubert"]).to(model.device)
 
     # Load the CustomTokenizer model
-    tokenizer = HubertTokenizer.load_from_checkpoint(model.config.LOCAL_MODEL_PATHS["hubert_tokenizer"]).to(
-        model.device
-    )  # Automatically uses
+    tokenizer = HubertTokenizer.load_from_checkpoint(model.config.LOCAL_MODEL_PATHS["hubert_tokenizer"], map_location=model.device)
     # semantic_tokens = model.text_to_semantic(
     #     text, max_gen_duration_s=seconds, top_k=50, top_p=0.95, temp=0.7
     # )  # not 100%