import os
import random
import sys
import torch
import torch.nn.functional as F
import torch.utils.data
from TTS.tts.models.xtts import load_audio
torch.set_num_threads(1)
def key_samples_by_col(samples, col):
"""Returns a dictionary of samples keyed by language."""
samples_by_col = {}
for sample in samples:
col_val = sample[col]
assert isinstance(col_val, str)
if col_val not in samples_by_col:
samples_by_col[col_val] = []
samples_by_col[col_val].append(sample)
return samples_by_col
def get_prompt_slice(gt_path, max_sample_length, min_sample_length, sample_rate, is_eval=False):
rel_clip = load_audio(gt_path, sample_rate)
# if eval uses a middle size sample when it is possible to be more reproducible
if is_eval:
sample_length = int((min_sample_length + max_sample_length) / 2)
else:
sample_length = random.randint(min_sample_length, max_sample_length)
gap = rel_clip.shape[-1] - sample_length
if gap < 0:
sample_length = rel_clip.shape[-1] // 2
gap = rel_clip.shape[-1] - sample_length
# if eval start always from the position 0 to be more reproducible
if is_eval:
rand_start = 0
else:
rand_start = random.randint(0, gap)
rand_end = rand_start + sample_length
rel_clip = rel_clip[:, rand_start:rand_end]
rel_clip = F.pad(rel_clip, pad=(0, max_sample_length - rel_clip.shape[-1]))
cond_idxs = [rand_start, rand_end]
return rel_clip, rel_clip.shape[-1], cond_idxs
class XTTSDataset(torch.utils.data.Dataset):
def __init__(self, config, samples, tokenizer, sample_rate, is_eval=False):
self.config = config
model_args = config.model_args
self.failed_samples = set()
self.debug_failures = model_args.debug_loading_failures
self.max_conditioning_length = model_args.max_conditioning_length
self.min_conditioning_length = model_args.min_conditioning_length
self.is_eval = is_eval
self.tokenizer = tokenizer
self.sample_rate = sample_rate
self.max_wav_len = model_args.max_wav_length
self.max_text_len = model_args.max_text_length
self.use_masking_gt_prompt_approach = model_args.gpt_use_masking_gt_prompt_approach
assert self.max_wav_len is not None and self.max_text_len is not None
self.samples = samples
if not is_eval:
random.seed(config.training_seed)
# random.shuffle(self.samples)
random.shuffle(self.samples)
# order by language
self.samples = key_samples_by_col(self.samples, "language")
print(" > Sampling by language:", self.samples.keys())
else:
# for evaluation load and check samples that are corrupted to ensures the reproducibility
self.check_eval_samples()
def check_eval_samples(self):
print(" > Filtering invalid eval samples!!")
new_samples = []
for sample in self.samples:
try:
tseq, _, wav, _, _, _ = self.load_item(sample)
except:
continue
# Basically, this audio file is nonexistent or too long to be supported by the dataset.
if (
wav is None
or (self.max_wav_len is not None and wav.shape[-1] > self.max_wav_len)
or (self.max_text_len is not None and tseq.shape[0] > self.max_text_len)
):
continue
new_samples.append(sample)
self.samples = new_samples
print(" > Total eval samples after filtering:", len(self.samples))
def get_text(self, text, lang):
tokens = self.tokenizer.encode(text, lang)
tokens = torch.IntTensor(tokens)
assert not torch.any(tokens == 1), f"UNK token found in {text} -> {self.tokenizer.decode(tokens)}"
# The stop token should always be sacred.
assert not torch.any(tokens == 0), f"Stop token found in {text}"
return tokens
def load_item(self, sample):
text = str(sample["text"])
tseq = self.get_text(text, sample["language"])
audiopath = sample["audio_file"]
wav = load_audio(audiopath, self.sample_rate)
if text is None or len(text.strip()) == 0:
raise ValueError
if wav is None or wav.shape[-1] < (0.5 * self.sample_rate):
# Ultra short clips are also useless (and can cause problems within some models).
raise ValueError
if self.use_masking_gt_prompt_approach:
# get a slice from GT to condition the model
cond, _, cond_idxs = get_prompt_slice(
audiopath, self.max_conditioning_length, self.min_conditioning_length, self.sample_rate, self.is_eval
)
# if use masking do not use cond_len
cond_len = torch.nan
else:
ref_sample = (
sample["reference_path"]
if "reference_path" in sample and sample["reference_path"] is not None
else audiopath
)
cond, cond_len, _ = get_prompt_slice(
ref_sample, self.max_conditioning_length, self.min_conditioning_length, self.sample_rate, self.is_eval
)
# if do not use masking use cond_len
cond_idxs = torch.nan
return tseq, audiopath, wav, cond, cond_len, cond_idxs
def __getitem__(self, index):
if self.is_eval:
sample = self.samples[index]
sample_id = str(index)
else:
# select a random language
lang = random.choice(list(self.samples.keys()))
# select random sample
index = random.randint(0, len(self.samples[lang]) - 1)
sample = self.samples[lang][index]
# a unique id for each sampel to deal with fails
sample_id = lang + "_" + str(index)
# ignore samples that we already know that is not valid ones
if sample_id in self.failed_samples:
if self.debug_failures:
print(f"Ignoring sample {sample['audio_file']} because it was already ignored before !!")
# call get item again to get other sample
return self[1]
# try to load the sample, if fails added it to the failed samples list
try:
tseq, audiopath, wav, cond, cond_len, cond_idxs = self.load_item(sample)
except:
if self.debug_failures:
print(f"error loading {sample['audio_file']} {sys.exc_info()}")
self.failed_samples.add(sample_id)
return self[1]
# check if the audio and text size limits and if it out of the limits, added it failed_samples
if (
wav is None
or (self.max_wav_len is not None and wav.shape[-1] > self.max_wav_len)
or (self.max_text_len is not None and tseq.shape[0] > self.max_text_len)
):
# Basically, this audio file is nonexistent or too long to be supported by the dataset.
# It's hard to handle this situation properly. Best bet is to return the a random valid token and skew the dataset somewhat as a result.
if self.debug_failures and wav is not None and tseq is not None:
print(
f"error loading {sample['audio_file']}: ranges are out of bounds; {wav.shape[-1]}, {tseq.shape[0]}"
)
self.failed_samples.add(sample_id)
return self[1]
res = {
# 'real_text': text,
"text": tseq,
"text_lengths": torch.tensor(tseq.shape[0], dtype=torch.long),
"wav": wav,
"wav_lengths": torch.tensor(wav.shape[-1], dtype=torch.long),
"filenames": audiopath,
"conditioning": cond.unsqueeze(1),
"cond_lens": torch.tensor(cond_len, dtype=torch.long)
if cond_len is not torch.nan
else torch.tensor([cond_len]),
"cond_idxs": torch.tensor(cond_idxs) if cond_idxs is not torch.nan else torch.tensor([cond_idxs]),
}
return res
def __len__(self):
if self.is_eval:
return len(self.samples)
return sum([len(v) for v in self.samples.values()])
def collate_fn(self, batch):
# convert list of dicts to dict of lists
B = len(batch)
batch = {k: [dic[k] for dic in batch] for k in batch[0]}
# stack for features that already have the same shape
batch["wav_lengths"] = torch.stack(batch["wav_lengths"])
batch["text_lengths"] = torch.stack(batch["text_lengths"])
batch["conditioning"] = torch.stack(batch["conditioning"])
batch["cond_lens"] = torch.stack(batch["cond_lens"])
batch["cond_idxs"] = torch.stack(batch["cond_idxs"])
if torch.any(batch["cond_idxs"].isnan()):
batch["cond_idxs"] = None
if torch.any(batch["cond_lens"].isnan()):
batch["cond_lens"] = None
max_text_len = batch["text_lengths"].max()
max_wav_len = batch["wav_lengths"].max()
# create padding tensors
text_padded = torch.IntTensor(B, max_text_len)
wav_padded = torch.FloatTensor(B, 1, max_wav_len)
# initialize tensors for zero padding
text_padded = text_padded.zero_()
wav_padded = wav_padded.zero_()
for i in range(B):
text = batch["text"][i]
text_padded[i, : batch["text_lengths"][i]] = torch.IntTensor(text)
wav = batch["wav"][i]
wav_padded[i, :, : batch["wav_lengths"][i]] = torch.FloatTensor(wav)
batch["wav"] = wav_padded
batch["padded_text"] = text_padded
return batch