Fix encoding header in meldataset.py

meldataset.py

@@ -7,7 +7,10 @@ import soundfile as sf
 import librosa
 
 import torch
-import torchaudio
+try:
+    import torchaudio
+except ImportError:
+    torchaudio = None
 import torch.utils.data
 import torch.distributed as dist
 from multiprocessing import Pool
@@ -18,115 +21,194 @@ logger.setLevel(logging.DEBUG)
 
 import pandas as pd
 
-class TextCleaner:
-    def __init__(self, symbol_dict, debug=True):
-        self.word_index_dictionary = symbol_dict
-        self.debug = debug
-    def __call__(self, text):
-        indexes = []
-        for char in text:
-            try:
-                indexes.append(self.word_index_dictionary[char])
-            except KeyError as e:
-                if self.debug:
-                    print("\nWARNING UNKNOWN IPA CHARACTERS/LETTERS: ", char)
-                    print("To ignore set 'debug' to false in the config")
-                continue
-        return indexes
-
-np.random.seed(1)
-random.seed(1)
+# class TextCleaner:
+#     def __init__(self, symbol_dict, debug=True):
+#         self.word_index_dictionary = symbol_dict
+#         self.debug = debug
+#     def __call__(self, text):
+#         indexes = []
+#         for char in text:
+#             try:
+#                 indexes.append(self.word_index_dictionary[char])
+#             except KeyError as e:
+#                 if self.debug:
+#                     print("\nWARNING UNKNOWN IPA CHARACTERS/LETTERS: ", char)
+#                     print("To ignore set 'debug' to false in the config")
+#                 continue
+#         return indexes
+
+
+
 SPECT_PARAMS = {
     "n_fft": 2048,
     "win_length": 1200,
-    "hop_length": 300
+    "hop_length": 300,
 }
+
+# Dùng đầy đủ params cho MelSpectrogram (tránh thiếu n_fft/win/hop)
 MEL_PARAMS = {
     "n_mels": 80,
+    "n_fft": 2048,
+    "win_length": 1200,
+    "hop_length": 300,
 }
 
-to_mel = torchaudio.transforms.MelSpectrogram(
-    n_mels=80, n_fft=2048, win_length=1200, hop_length=300)
 mean, std = -4, 4
 
-def preprocess(wave):
+
+
+# Cache MelSpectrogram theo sample_rate
+_MEL_CACHE = {}
+
+
+def _require_torchaudio(context: str) -> None:
+    if torchaudio is None:
+        raise RuntimeError(
+            f"torchaudio is required for {context} but is not installed in this environment. "
+            "For HF Spaces inference, you should not instantiate FilePathDataset / mel extraction."
+        )
+
+
+def get_mel_transform(sample_rate: int = 24000):
+    _require_torchaudio("mel extraction")
+    if sample_rate not in _MEL_CACHE:
+        _MEL_CACHE[sample_rate] = torchaudio.transforms.MelSpectrogram(
+            sample_rate=sample_rate,
+            n_mels=MEL_PARAMS["n_mels"],
+            n_fft=MEL_PARAMS["n_fft"],
+            win_length=MEL_PARAMS["win_length"],
+            hop_length=MEL_PARAMS["hop_length"],
+        )
+    return _MEL_CACHE[sample_rate]
+
+
+def preprocess(wave: np.ndarray, sample_rate: int = 24000):
+    """
+    wave: 1D numpy float array
+    return: mel tensor shape (1, n_mels, T)
+    """
+    _require_torchaudio("preprocess()")
+    if wave.ndim != 1:
+        wave = np.asarray(wave).squeeze()
     wave_tensor = torch.from_numpy(wave).float()
-    mel_tensor = to_mel(wave_tensor)
-    mel_tensor = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - mean) / std
-    return mel_tensor
+
+    to_mel = get_mel_transform(sample_rate)
+    mel = to_mel(wave_tensor)  # (n_mels, T)
+    mel = (torch.log(mel + 1e-5) - mean) / std
+    return mel.unsqueeze(0)  # (1, n_mels, T)
+
+
+class TextCleaner:
+    """
+    Minimal TextCleaner: map token -> id based on symbol_dict.
+    - Nếu input text có dấu cách: split theo space (phù hợp IPA tokenization)
+    - Nếu không có space: tách theo ký tự
+    """
+    def __init__(self, symbol_dict, debug=True):
+        self.symbol_dict = symbol_dict
+        self.debug = debug
+
+    def __call__(self, text: str):
+        text = (text or "").strip()
+
+        # IPA/token list thường được tách bằng space
+        if " " in text:
+            tokens = [t for t in text.split(" ") if t != ""]
+        else:
+            tokens = list(text)
+
+        ids = []
+        missing = []
+        for t in tokens:
+            if t in self.symbol_dict:
+                ids.append(self.symbol_dict[t])
+            else:
+                missing.append(t)
+
+        if self.debug and missing:
+            # In tối đa 30 token thiếu để tránh spam log
+            print(f"[TextCleaner] missing {len(missing)} symbols. sample={missing[:30]}")
+
+        return ids
+
 
 class FilePathDataset(torch.utils.data.Dataset):
-    def __init__(self,
-                 data_list,
-                 root_path,
-                 symbol_dict,
-                 sr=24000,
-                 data_augmentation=False,
-                 validation=False,
-                 debug=True
-                 ):
-
-        _data_list = [l.strip().split('|') for l in data_list]
-        self.data_list = _data_list #[data if len(data) == 3 else (*data, 0) for data in _data_list] #append speakerid=0 for all
+    def __init__(
+        self,
+        data_list,
+        root_path,
+        symbol_dict,
+        sr=24000,
+        data_augmentation=False,
+        validation=False,
+        debug=True,
+    ):
+        _require_torchaudio("FilePathDataset (training dataloader)")
+
+        _data_list = [l.strip().split("|") for l in data_list]
+        self.data_list = _data_list  # [wav_path, text] (hoặc thêm speaker_id tuỳ bạn)
         self.text_cleaner = TextCleaner(symbol_dict, debug)
         self.sr = sr
 
         self.df = pd.DataFrame(self.data_list)
 
-        self.to_melspec = torchaudio.transforms.MelSpectrogram(**MEL_PARAMS)
+        # training-only: mel transform
+        self.to_melspec = get_mel_transform(self.sr)
 
         self.mean, self.std = -4, 4
         self.data_augmentation = data_augmentation and (not validation)
         self.max_mel_length = 192
-        
         self.root_path = root_path
 
     def __len__(self):
         return len(self.data_list)
 
-    def __getitem__(self, idx):        
+    def __getitem__(self, idx):
         data = self.data_list[idx]
         path = data[0]
-        
+
         wave, text_tensor = self._load_tensor(data)
-        
-        mel_tensor = preprocess(wave).squeeze()
-        
-        acoustic_feature = mel_tensor.squeeze()
+
+        mel_tensor = preprocess(wave, sample_rate=self.sr).squeeze()  # (n_mels, T)
+
+        acoustic_feature = mel_tensor
         length_feature = acoustic_feature.size(1)
-        acoustic_feature = acoustic_feature[:, :(length_feature - length_feature % 2)]
-        
+        acoustic_feature = acoustic_feature[:, : (length_feature - length_feature % 2)]
+
         return acoustic_feature, text_tensor, path, wave
 
     def _load_tensor(self, data):
-        wave_path, text = data
+        # data có thể là [wave_path, text] hoặc [wave_path, text, speaker_id]
+        wave_path = data[0]
+        text = data[1]
+
         wave, sr = sf.read(osp.join(self.root_path, wave_path))
-        if wave.shape[-1] == 2:
+        if isinstance(wave, np.ndarray) and wave.ndim == 2 and wave.shape[-1] == 2:
             wave = wave[:, 0].squeeze()
-        if sr != 24000:
-            wave = librosa.resample(wave, orig_sr=sr, target_sr=24000)
-            print(wave_path, sr)
-        
-        # Adding half a second padding.
-        wave = np.concatenate([np.zeros([12000]), wave, np.zeros([12000])], axis=0) 
-        
-        text = self.text_cleaner(text)
-        
-        text.insert(0, 0)
-        text.append(0)
-        
-        text = torch.LongTensor(text)
 
-        return wave, text
+        if sr != self.sr:
+            wave = librosa.resample(wave, orig_sr=sr, target_sr=self.sr)
+
+        # padding 0.5s mỗi bên (24000 * 0.5 = 12000)
+        wave = np.concatenate([np.zeros([12000]), wave, np.zeros([12000])], axis=0)
+
+        text_ids = self.text_cleaner(text)
+
+        # BOS/EOS = 0 như code gốc của bạn
+        text_ids.insert(0, 0)
+        text_ids.append(0)
+
+        text_tensor = torch.LongTensor(text_ids)
+        return wave, text_tensor
 
     def _load_data(self, data):
         wave, text_tensor = self._load_tensor(data)
-        mel_tensor = preprocess(wave).squeeze()
+        mel_tensor = preprocess(wave, sample_rate=self.sr).squeeze()
 
         mel_length = mel_tensor.size(1)
         if mel_length > self.max_mel_length:
             random_start = np.random.randint(0, mel_length - self.max_mel_length)
-            mel_tensor = mel_tensor[:, random_start:random_start + self.max_mel_length]
+            mel_tensor = mel_tensor[:, random_start : random_start + self.max_mel_length]
 
         return mel_tensor