import math
from pathlib import Path
import random
from typing import List, Tuple
import torch
from torch import Tensor
from torch.nn import functional as F
from torch.utils.data import DataLoader, Dataset
from models.config import HifiGanPretrainingConfig
from .hifi_libri_dataset import NUM_JOBS, HifiLibriDataset
class HifiGanDataset(Dataset):
r"""A PyTorch Dataset for the HiFi-GAN model.
Args:
lang (str, optional): The language of the dataset. Defaults to "en".
root (str, optional): The root directory of the dataset. Defaults to "datasets_cache".
sampling_rate (int, optional): The sampling rate of the audio. Defaults to 44100.
hifitts_path (str, optional): The path to the HiFiTTS dataset. Defaults to "hifitts".
hifi_cutset_file_name (str, optional): The file name of the HiFiTTS cutset. Defaults to "hifi.json.gz".
libritts_path (str, optional): The path to the LibriTTS dataset. Defaults to "librittsr".
libritts_cutset_file_name (str, optional): The file name of the LibriTTS cutset. Defaults to "libri.json.gz".
libritts_subsets (Union[List[str], str], optional): The subsets of the LibriTTS dataset to use. Defaults to "all".
cache (bool, optional): Whether to cache the dataset. Defaults to False.
cache_dir (str, optional): The directory to cache the dataset in. Defaults to "/dev/shm".
num_jobs (int, optional): The number of jobs to use for preparing the dataset. Defaults to NUM_JOBS.
"""
def __init__(
self,
lang: str = "en",
root: str = "datasets_cache",
sampling_rate: int = 44100,
hifitts_path: str = "hifitts",
hifi_cutset_file_name: str = "hifi.json.gz",
libritts_path: str = "librittsr",
libritts_cutset_file_name: str = "libri.json.gz",
libritts_subsets: List[str] | str = "all",
cache: bool = False,
cache_dir: str = "/dev/shm",
num_jobs: int = NUM_JOBS,
):
self.cache = cache
self.cache_dir = Path(cache_dir) / "cache-hifigan-dataset"
self.pretraining_config = HifiGanPretrainingConfig()
self.dataset = HifiLibriDataset(
lang=lang,
root=root,
sampling_rate=sampling_rate,
hifitts_path=hifitts_path,
hifi_cutset_file_name=hifi_cutset_file_name,
libritts_path=libritts_path,
libritts_cutset_file_name=libritts_cutset_file_name,
libritts_subsets=libritts_subsets,
num_jobs=num_jobs,
max_seconds=25.0, # To be sure that all the audio files from the dataset will be used
include_libri=False, # Exclude LibriTTS dataset
)
self.segment_size = self.pretraining_config.segment_size
self.hop_size = self.dataset.preprocess_config.stft.hop_length
def __len__(self):
r"""Get the number of items in the dataset.
Returns:
int: The number of items in the dataset.
"""
return len(self.dataset)
def get_cache_subdir_path(self, idx: int) -> Path:
r"""Calculate the path to the cache subdirectory.
Args:
idx (int): The index of the cache subdirectory.
Returns:
Path: The path to the cache subdirectory.
"""
return self.cache_dir / str(((idx // 1000) + 1) * 1000)
def get_cache_file_path(self, idx: int) -> Path:
r"""Calculate the path to the cache file.
Args:
idx (int): The index of the cache file.
Returns:
Path: The path to the cache file.
"""
return self.get_cache_subdir_path(idx) / f"{idx}.pt"
def __getitem__(self, idx: int) -> Tuple[str, Tensor, Tensor]:
r"""Get an item from the dataset.
If caching is enabled and the item is in the cache, the cached item is returned.
Otherwise, the item is loaded from the dataset, preprocessed, and returned.
Args:
idx (int): The index of the item in the dataset.
Returns:
Tuple[str, Tensor, Tensor]: The ID of the item, the audio waveform, and the mel spectrogram.
"""
cache_file = self.get_cache_file_path(idx)
if self.cache and cache_file.exists():
cached_data: Tuple[str, Tensor, Tensor] = torch.load(cache_file)
return cached_data
item = self.dataset[idx]
frames_per_seg = math.ceil(self.segment_size / self.hop_size)
audio = item.wav
mel = item.mel.unsqueeze(0)
if audio.size(1) >= self.segment_size:
mel_start = random.randint(0, mel.size(2) - frames_per_seg - 1) # noqa: S311
mel = mel[:, :, mel_start : mel_start + frames_per_seg]
audio = audio[
:,
mel_start * self.hop_size : (mel_start + frames_per_seg)
* self.hop_size,
]
else:
mel = F.pad(
mel,
(0, frames_per_seg - mel.size(2)),
"constant",
)
audio = F.pad(
audio,
(0, self.segment_size - audio.size(1)),
"constant",
)
result = (item.id, audio, mel.squeeze(0))
if self.cache:
# Create the cache subdirectory if it doesn't exist
Path.mkdir(
self.get_cache_subdir_path(idx),
parents=True,
exist_ok=True,
)
# Save the preprocessed data to the cache
torch.save(result, cache_file)
return result
def __iter__(self):
r"""Method makes the class iterable. It iterates over the `_walker` attribute
and for each item, it gets the corresponding item from the dataset using the
`__getitem__` method.
Yields:
The item from the dataset corresponding to the current item in `_walker`.
"""
for item in range(self.__len__()):
yield self.__getitem__(item)
def train_dataloader(
lang: str = "en",
root: str = "datasets_cache",
sampling_rate: int = 44100,
hifitts_path: str = "hifitts",
hifi_cutset_file_name: str = "hifi.json.gz",
libritts_path: str = "librittsr",
libritts_cutset_file_name: str = "libri.json.gz",
libritts_subsets: List[str] | str = "all",
cache: bool = False,
cache_dir: str = "/dev/shm",
num_jobs: int = NUM_JOBS,
num_workers: int = 0,
shuffle: bool = False,
batch_size: int = 5,
pin_memory: bool = True,
drop_last: bool = True,
) -> DataLoader:
r"""Create a DataLoader for the training data.
Args:
lang (str, optional): The language of the dataset. Defaults to "en".
root (str, optional): The root directory of the dataset. Defaults to "datasets_cache".
sampling_rate (int, optional): The sampling rate of the audio. Defaults to 44100.
hifitts_path (str, optional): The path to the HiFiTTS dataset. Defaults to "hifitts".
hifi_cutset_file_name (str, optional): The file name of the HiFiTTS cutset. Defaults to "hifi.json.gz".
libritts_path (str, optional): The path to the LibriTTS dataset. Defaults to "librittsr".
libritts_cutset_file_name (str, optional): The file name of the LibriTTS cutset. Defaults to "libri.json.gz".
libritts_subsets (Union[List[str], str], optional): The subsets of the LibriTTS dataset to use. Defaults to "all".
cache (bool, optional): Whether to cache the dataset. Defaults to False.
cache_dir (str, optional): The directory to cache the dataset in. Defaults to "/dev/shm".
num_jobs (int, optional): The number of jobs to use for preparing the dataset. Defaults to NUM_JOBS.
num_workers (int, optional): The number of worker processes to use for loading the data. Defaults to 0.
shuffle (bool, optional): Whether to shuffle the data. Defaults to False.
batch_size (int, optional): The batch size. Defaults to 5.
pin_memory (bool, optional): Whether to pin memory. Defaults to True.
drop_last (bool, optional): Whether to drop the last incomplete batch. Defaults to True.
num_gpus (int, optional): The number of GPUs to use. Defaults to 1.
Returns:
DataLoader: A DataLoader for the training data.
"""
trainset = HifiGanDataset(
lang=lang,
root=root,
sampling_rate=sampling_rate,
hifitts_path=hifitts_path,
hifi_cutset_file_name=hifi_cutset_file_name,
libritts_path=libritts_path,
libritts_cutset_file_name=libritts_cutset_file_name,
libritts_subsets=libritts_subsets,
cache=cache,
cache_dir=cache_dir,
num_jobs=num_jobs,
)
train_loader = DataLoader(
trainset,
num_workers=num_workers,
shuffle=shuffle,
batch_size=batch_size,
pin_memory=pin_memory,
drop_last=drop_last,
)
return train_loader