Update yarngpt_utils.py

yarngpt_utils.py

@@ -1,6 +1,7 @@
 # yarngpt_utils.py
 import torch
 import torchaudio
+import re
 from outetts.wav_tokenizer.decoder import WavTokenizer
 from transformers import AutoTokenizer
 
@@ -8,11 +9,14 @@ class AudioTokenizer:
     def __init__(self, hf_path, wav_tokenizer_model_path, wav_tokenizer_config_path):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.tokenizer = AutoTokenizer.from_pretrained(hf_path)
+        
+        # Correct initialization for WavTokenizer
         self.wav_tokenizer = WavTokenizer(
-            checkpoint_path=wav_tokenizer_model_path,
+            model_path=wav_tokenizer_model_path,
             config_path=wav_tokenizer_config_path,
             device=self.device
         )
+        
         self.speakers = ["idera", "emma", "jude", "osagie", "tayo", "zainab", 
                          "joke", "regina", "remi", "umar", "chinenye"]