stable_codec/data/dataset.py

import importlib
import random
import torch
import torchaudio
import webdataset as wds

from typing import List
from torchaudio import transforms as T

from stable_audio_tools.data.dataset import (
    S3DatasetConfig, LocalWebDatasetConfig, log_and_continue, audio_decoder, npy_decoder,
    is_valid_sample, collation_fn, AUDIO_KEYS, remove_long_silence,
)
from stable_audio_tools.data.utils import (
    Stereo, Mono, PhaseFlipper, PadCrop_Normalized_T, VolumeNorm,
)

from .Text2Phone.Text2PhoneTokenizer import Text2PhoneTokenizer

class WebDatasetDataLoader():
    def __init__(
        self,
        datasets: List[S3DatasetConfig],
        batch_size,
        sample_size,
        sample_rate=48000,
        num_workers=8,
        epoch_steps=1000,
        random_crop=True,
        force_channels="stereo",
        augment_phase=True,
        remove_silence=True,
        silence_threshold=[0.01, 0.5],
        max_silence_duration=0.2,
        volume_norm=False,
        volume_norm_param=(-16, 2),
        pre_encoded=False,
        resampled_shards=True,
        force_align_text=False,
        **data_loader_kwargs
    ):

        self.datasets = datasets

        self.sample_size = sample_size
        self.sample_rate = sample_rate
        self.random_crop = random_crop
        self.force_channels = force_channels
        self.augment_phase = augment_phase
        self.pre_encoded = pre_encoded
        self.volume_norm = volume_norm
        self.volume_norm_param = volume_norm_param
        self.remove_silence = remove_silence
        self.silence_threshold = silence_threshold
        self.max_silence_duration = max_silence_duration

        self.force_align_text = force_align_text
        self.phonemizer = Text2PhoneTokenizer()

        urls = [dataset.load_data_urls() for dataset in datasets]

        # Flatten the list of lists of URLs
        urls = [url for dataset_urls in urls for url in dataset_urls]

        # Shuffle the urls
        random.shuffle(urls)

        self.dataset = wds.DataPipeline(
            wds.ResampledShards(urls) if resampled_shards else wds.SimpleShardList(urls),
            wds.tarfile_to_samples(handler=log_and_continue),
            wds.decode(audio_decoder, handler=log_and_continue) if not self.pre_encoded else wds.decode(npy_decoder, handler=log_and_continue),
            wds.map(self.wds_preprocess, handler=log_and_continue),
            wds.select(is_valid_sample),
            wds.to_tuple("audio", "json", handler=log_and_continue),
            #wds.shuffle(bufsize=1000, initial=5000),
            wds.batched(batch_size, partial=False, collation_fn=collation_fn),
        )

        if resampled_shards:
            self.dataset = self.dataset.with_epoch(epoch_steps//num_workers if num_workers > 0 else epoch_steps)

        self.data_loader = wds.WebLoader(self.dataset, num_workers=num_workers, **data_loader_kwargs)

    def wds_preprocess(self, sample):

        if self.pre_encoded:
            audio = torch.from_numpy(sample["npy"])
            del sample["npy"]
            sample["__pre_encoded__"] = True
            sample["json"]["padding_mask"] = torch.tensor(sample["json"]["padding_mask"])
        else:
            found_key, rewrite_key = '', ''
            for k, v in sample.items():  # print the all entries in dict
                for akey in AUDIO_KEYS:
                    if k.endswith(akey):
                        # to rename long/weird key with its simpler counterpart
                        found_key, rewrite_key = k, akey
                        break
                if '' != found_key:
                    break
            if '' == found_key:  # got no audio!
                return None  # try returning None to tell WebDataset to skip this one

            audio, in_sr = sample[found_key]
            if in_sr != self.sample_rate:
                resample_tf = T.Resample(in_sr, self.sample_rate)
                audio = resample_tf(audio)

            # Replace the long silence by the short for the mono audios
            if audio.shape[0] == 1 and self.remove_silence:
                audio = remove_long_silence(audio, self.sample_rate, self.silence_threshold, self.max_silence_duration)

            original_length = audio.shape[-1]

            if self.sample_size is not None:
                # Pad/crop and get the relative timestamp
                pad_crop = PadCrop_Normalized_T(
                    self.sample_size, randomize=self.random_crop, sample_rate=self.sample_rate)
                audio, t_start, t_end, seconds_start, seconds_total, padding_mask = pad_crop(
                    audio)
                sample["json"]["seconds_start"] = seconds_start
                sample["json"]["seconds_total"] = seconds_total
                sample["json"]["padding_mask"] = padding_mask
            else:
                t_start, t_end = 0, 1

            start_time = (original_length * t_start) / self.sample_rate
            end_time = (original_length * t_end) / self.sample_rate

            # Check if audio is length zero, initialize to a single zero if so
            if audio.shape[-1] == 0:
                audio = torch.zeros(1, 1)

            # Make the audio stereo and augment by randomly inverting phase
            augs = torch.nn.Sequential(
                Stereo() if self.force_channels == "stereo" else torch.nn.Identity(),
                Mono() if self.force_channels == "mono" else torch.nn.Identity(),
                VolumeNorm(self.volume_norm_param, self.sample_rate) if self.volume_norm else torch.nn.Identity(),
                PhaseFlipper() if self.augment_phase else torch.nn.Identity()
            )

            audio = augs(audio)

            sample["json"]["timestamps"] = (t_start, t_end)

            if found_key != rewrite_key:   # rename long/weird key with its simpler counterpart
                del sample[found_key]

        if "text" in sample["json"]:
            sample["json"]["prompt"] = sample["json"]["text"]

        # Check for custom metadata functions
        for dataset in self.datasets:
            if dataset.custom_metadata_fn is None:
                continue

            if dataset.path in sample["__url__"]:
                custom_metadata = dataset.custom_metadata_fn(sample["json"], audio)
                sample["json"].update(custom_metadata)

        sample["audio"] = audio
        # Add audio to the metadata as well for conditioning
        sample["json"]["audio"] = audio

        if self.force_align_text and self.sample_size is not None:
            # Chunk the original transcriptions according to (start_time, end_time)
            chunked_text_list = []
            for entry in sample["json"]['force_aligned_text']['transcript']:
                word_start = entry['start']
                word_end = entry['end']
                # Check if the word's start or end time falls within the time range
                if (word_start >= start_time and word_start <= end_time) or (word_end >= start_time and word_end <= end_time):
                    chunked_text_list.append(entry['word'])

            chunked_text =  ' '.join(chunked_text_list)
            chunked_phone = self.phonemizer.tokenize(chunked_text)

            sample["json"]["phone"] = chunked_phone
            sample["json"]["aligned_text"] = chunked_text

        return sample

def create_dataloader_from_config(dataset_config, batch_size, sample_size, sample_rate, audio_channels=2, num_workers=4, shuffle = True):
    dataset_type = dataset_config.get("dataset_type", None)
    assert dataset_type is not None, "Dataset type must be specified in dataset config"
    assert dataset_type in ("s3", "wds")

    force_channels = "mono" if audio_channels == 1 else "stereo"

    wds_configs = []
    for wds_config in dataset_config["datasets"]:
        custom_metadata_fn = None
        custom_metadata_module_path = wds_config.get("custom_metadata_module", None)

        if custom_metadata_module_path is not None:
            spec = importlib.util.spec_from_file_location(
                "metadata_module", custom_metadata_module_path)
            metadata_module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(metadata_module)
            custom_metadata_fn = metadata_module.get_custom_metadata

        if "s3_path" in wds_config:
            wds_configs.append(S3DatasetConfig(
                id=wds_config["id"],
                s3_path=wds_config["s3_path"],
                custom_metadata_fn=custom_metadata_fn,
                profile=wds_config.get("profile", None),
            ))
        elif "path" in wds_config:
            wds_configs.append(LocalWebDatasetConfig(
                id=wds_config["id"],
                path=wds_config["path"],
                custom_metadata_fn=custom_metadata_fn
            ))

    return WebDatasetDataLoader(
        wds_configs,
        sample_rate=sample_rate,
        sample_size=sample_size,
        batch_size=batch_size,
        remove_silence=dataset_config.get("remove_silence", False),
        silence_threshold=dataset_config.get("silence_threshold", [0.01, 0.5]),
        max_silence_duration=dataset_config.get("max_silence_duration", 0.25),
        random_crop=dataset_config.get("random_crop", True),
        volume_norm=dataset_config.get("volume_norm", False),
        volume_norm_param=dataset_config.get("volume_norm_param", [-16, 2]),
        num_workers=num_workers,
        persistent_workers=True,
        pin_memory=True,
        force_channels=force_channels,
        epoch_steps=dataset_config.get("epoch_steps", 2000),
        pre_encoded=dataset_config.get("pre_encoded", False),
        resampled_shards=dataset_config.get("resampled_shards", True),
        force_align_text=dataset_config.get("force_align_text", False)
    ).data_loader