xinference 0.14.2__py3-none-any.whl → 0.14.4__py3-none-any.whl

xinference/thirdparty/matcha/cli.py

@@ -0,0 +1,419 @@
+import argparse
+import datetime as dt
+import os
+import warnings
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import soundfile as sf
+import torch
+
+from matcha.hifigan.config import v1
+from matcha.hifigan.denoiser import Denoiser
+from matcha.hifigan.env import AttrDict
+from matcha.hifigan.models import Generator as HiFiGAN
+from matcha.models.matcha_tts import MatchaTTS
+from matcha.text import sequence_to_text, text_to_sequence
+from matcha.utils.utils import assert_model_downloaded, get_user_data_dir, intersperse
+
+MATCHA_URLS = {
+    "matcha_ljspeech": "https://github.com/shivammehta25/Matcha-TTS-checkpoints/releases/download/v1.0/matcha_ljspeech.ckpt",
+    "matcha_vctk": "https://github.com/shivammehta25/Matcha-TTS-checkpoints/releases/download/v1.0/matcha_vctk.ckpt",
+}
+
+VOCODER_URLS = {
+    "hifigan_T2_v1": "https://github.com/shivammehta25/Matcha-TTS-checkpoints/releases/download/v1.0/generator_v1",  # Old url: https://drive.google.com/file/d/14NENd4equCBLyyCSke114Mv6YR_j_uFs/view?usp=drive_link
+    "hifigan_univ_v1": "https://github.com/shivammehta25/Matcha-TTS-checkpoints/releases/download/v1.0/g_02500000",  # Old url: https://drive.google.com/file/d/1qpgI41wNXFcH-iKq1Y42JlBC9j0je8PW/view?usp=drive_link
+}
+
+MULTISPEAKER_MODEL = {
+    "matcha_vctk": {"vocoder": "hifigan_univ_v1", "speaking_rate": 0.85, "spk": 0, "spk_range": (0, 107)}
+}
+
+SINGLESPEAKER_MODEL = {"matcha_ljspeech": {"vocoder": "hifigan_T2_v1", "speaking_rate": 0.95, "spk": None}}
+
+
+def plot_spectrogram_to_numpy(spectrogram, filename):
+    fig, ax = plt.subplots(figsize=(12, 3))
+    im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
+    plt.colorbar(im, ax=ax)
+    plt.xlabel("Frames")
+    plt.ylabel("Channels")
+    plt.title("Synthesised Mel-Spectrogram")
+    fig.canvas.draw()
+    plt.savefig(filename)
+
+
+def process_text(i: int, text: str, device: torch.device):
+    print(f"[{i}] - Input text: {text}")
+    x = torch.tensor(
+        intersperse(text_to_sequence(text, ["english_cleaners2"])[0], 0),
+        dtype=torch.long,
+        device=device,
+    )[None]
+    x_lengths = torch.tensor([x.shape[-1]], dtype=torch.long, device=device)
+    x_phones = sequence_to_text(x.squeeze(0).tolist())
+    print(f"[{i}] - Phonetised text: {x_phones[1::2]}")
+
+    return {"x_orig": text, "x": x, "x_lengths": x_lengths, "x_phones": x_phones}
+
+
+def get_texts(args):
+    if args.text:
+        texts = [args.text]
+    else:
+        with open(args.file, encoding="utf-8") as f:
+            texts = f.readlines()
+    return texts
+
+
+def assert_required_models_available(args):
+    save_dir = get_user_data_dir()
+    if not hasattr(args, "checkpoint_path") and args.checkpoint_path is None:
+        model_path = args.checkpoint_path
+    else:
+        model_path = save_dir / f"{args.model}.ckpt"
+        assert_model_downloaded(model_path, MATCHA_URLS[args.model])
+
+    vocoder_path = save_dir / f"{args.vocoder}"
+    assert_model_downloaded(vocoder_path, VOCODER_URLS[args.vocoder])
+    return {"matcha": model_path, "vocoder": vocoder_path}
+
+
+def load_hifigan(checkpoint_path, device):
+    h = AttrDict(v1)
+    hifigan = HiFiGAN(h).to(device)
+    hifigan.load_state_dict(torch.load(checkpoint_path, map_location=device)["generator"])
+    _ = hifigan.eval()
+    hifigan.remove_weight_norm()
+    return hifigan
+
+
+def load_vocoder(vocoder_name, checkpoint_path, device):
+    print(f"[!] Loading {vocoder_name}!")
+    vocoder = None
+    if vocoder_name in ("hifigan_T2_v1", "hifigan_univ_v1"):
+        vocoder = load_hifigan(checkpoint_path, device)
+    else:
+        raise NotImplementedError(
+            f"Vocoder {vocoder_name} not implemented! define a load_<<vocoder_name>> method for it"
+        )
+
+    denoiser = Denoiser(vocoder, mode="zeros")
+    print(f"[+] {vocoder_name} loaded!")
+    return vocoder, denoiser
+
+
+def load_matcha(model_name, checkpoint_path, device):
+    print(f"[!] Loading {model_name}!")
+    model = MatchaTTS.load_from_checkpoint(checkpoint_path, map_location=device)
+    _ = model.eval()
+
+    print(f"[+] {model_name} loaded!")
+    return model
+
+
+def to_waveform(mel, vocoder, denoiser=None):
+    audio = vocoder(mel).clamp(-1, 1)
+    if denoiser is not None:
+        audio = denoiser(audio.squeeze(), strength=0.00025).cpu().squeeze()
+
+    return audio.cpu().squeeze()
+
+
+def save_to_folder(filename: str, output: dict, folder: str):
+    folder = Path(folder)
+    folder.mkdir(exist_ok=True, parents=True)
+    plot_spectrogram_to_numpy(np.array(output["mel"].squeeze().float().cpu()), f"{filename}.png")
+    np.save(folder / f"{filename}", output["mel"].cpu().numpy())
+    sf.write(folder / f"{filename}.wav", output["waveform"], 22050, "PCM_24")
+    return folder.resolve() / f"{filename}.wav"
+
+
+def validate_args(args):
+    assert (
+        args.text or args.file
+    ), "Either text or file must be provided Matcha-T(ea)TTS need sometext to whisk the waveforms."
+    assert args.temperature >= 0, "Sampling temperature cannot be negative"
+    assert args.steps > 0, "Number of ODE steps must be greater than 0"
+
+    if args.checkpoint_path is None:
+        # When using pretrained models
+        if args.model in SINGLESPEAKER_MODEL:
+            args = validate_args_for_single_speaker_model(args)
+
+        if args.model in MULTISPEAKER_MODEL:
+            args = validate_args_for_multispeaker_model(args)
+    else:
+        # When using a custom model
+        if args.vocoder != "hifigan_univ_v1":
+            warn_ = "[-] Using custom model checkpoint! I would suggest passing --vocoder hifigan_univ_v1, unless the custom model is trained on LJ Speech."
+            warnings.warn(warn_, UserWarning)
+        if args.speaking_rate is None:
+            args.speaking_rate = 1.0
+
+    if args.batched:
+        assert args.batch_size > 0, "Batch size must be greater than 0"
+    assert args.speaking_rate > 0, "Speaking rate must be greater than 0"
+
+    return args
+
+
+def validate_args_for_multispeaker_model(args):
+    if args.vocoder is not None:
+        if args.vocoder != MULTISPEAKER_MODEL[args.model]["vocoder"]:
+            warn_ = f"[-] Using {args.model} model! I would suggest passing --vocoder {MULTISPEAKER_MODEL[args.model]['vocoder']}"
+            warnings.warn(warn_, UserWarning)
+    else:
+        args.vocoder = MULTISPEAKER_MODEL[args.model]["vocoder"]
+
+    if args.speaking_rate is None:
+        args.speaking_rate = MULTISPEAKER_MODEL[args.model]["speaking_rate"]
+
+    spk_range = MULTISPEAKER_MODEL[args.model]["spk_range"]
+    if args.spk is not None:
+        assert (
+            args.spk >= spk_range[0] and args.spk <= spk_range[-1]
+        ), f"Speaker ID must be between {spk_range} for this model."
+    else:
+        available_spk_id = MULTISPEAKER_MODEL[args.model]["spk"]
+        warn_ = f"[!] Speaker ID not provided! Using speaker ID {available_spk_id}"
+        warnings.warn(warn_, UserWarning)
+        args.spk = available_spk_id
+
+    return args
+
+
+def validate_args_for_single_speaker_model(args):
+    if args.vocoder is not None:
+        if args.vocoder != SINGLESPEAKER_MODEL[args.model]["vocoder"]:
+            warn_ = f"[-] Using {args.model} model! I would suggest passing --vocoder {SINGLESPEAKER_MODEL[args.model]['vocoder']}"
+            warnings.warn(warn_, UserWarning)
+    else:
+        args.vocoder = SINGLESPEAKER_MODEL[args.model]["vocoder"]
+
+    if args.speaking_rate is None:
+        args.speaking_rate = SINGLESPEAKER_MODEL[args.model]["speaking_rate"]
+
+    if args.spk != SINGLESPEAKER_MODEL[args.model]["spk"]:
+        warn_ = f"[-] Ignoring speaker id {args.spk} for {args.model}"
+        warnings.warn(warn_, UserWarning)
+        args.spk = SINGLESPEAKER_MODEL[args.model]["spk"]
+
+    return args
+
+
+@torch.inference_mode()
+def cli():
+    parser = argparse.ArgumentParser(
+        description=" 🍵 Matcha-TTS: A fast TTS architecture with conditional flow matching"
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="matcha_ljspeech",
+        help="Model to use",
+        choices=MATCHA_URLS.keys(),
+    )
+
+    parser.add_argument(
+        "--checkpoint_path",
+        type=str,
+        default=None,
+        help="Path to the custom model checkpoint",
+    )
+
+    parser.add_argument(
+        "--vocoder",
+        type=str,
+        default=None,
+        help="Vocoder to use (default: will use the one suggested with the pretrained model))",
+        choices=VOCODER_URLS.keys(),
+    )
+    parser.add_argument("--text", type=str, default=None, help="Text to synthesize")
+    parser.add_argument("--file", type=str, default=None, help="Text file to synthesize")
+    parser.add_argument("--spk", type=int, default=None, help="Speaker ID")
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=0.667,
+        help="Variance of the x0 noise (default: 0.667)",
+    )
+    parser.add_argument(
+        "--speaking_rate",
+        type=float,
+        default=None,
+        help="change the speaking rate, a higher value means slower speaking rate (default: 1.0)",
+    )
+    parser.add_argument("--steps", type=int, default=10, help="Number of ODE steps  (default: 10)")
+    parser.add_argument("--cpu", action="store_true", help="Use CPU for inference (default: use GPU if available)")
+    parser.add_argument(
+        "--denoiser_strength",
+        type=float,
+        default=0.00025,
+        help="Strength of the vocoder bias denoiser (default: 0.00025)",
+    )
+    parser.add_argument(
+        "--output_folder",
+        type=str,
+        default=os.getcwd(),
+        help="Output folder to save results (default: current dir)",
+    )
+    parser.add_argument("--batched", action="store_true", help="Batched inference (default: False)")
+    parser.add_argument(
+        "--batch_size", type=int, default=32, help="Batch size only useful when --batched (default: 32)"
+    )
+
+    args = parser.parse_args()
+
+    args = validate_args(args)
+    device = get_device(args)
+    print_config(args)
+    paths = assert_required_models_available(args)
+
+    if args.checkpoint_path is not None:
+        print(f"[🍵] Loading custom model from {args.checkpoint_path}")
+        paths["matcha"] = args.checkpoint_path
+        args.model = "custom_model"
+
+    model = load_matcha(args.model, paths["matcha"], device)
+    vocoder, denoiser = load_vocoder(args.vocoder, paths["vocoder"], device)
+
+    texts = get_texts(args)
+
+    spk = torch.tensor([args.spk], device=device, dtype=torch.long) if args.spk is not None else None
+    if len(texts) == 1 or not args.batched:
+        unbatched_synthesis(args, device, model, vocoder, denoiser, texts, spk)
+    else:
+        batched_synthesis(args, device, model, vocoder, denoiser, texts, spk)
+
+
+class BatchedSynthesisDataset(torch.utils.data.Dataset):
+    def __init__(self, processed_texts):
+        self.processed_texts = processed_texts
+
+    def __len__(self):
+        return len(self.processed_texts)
+
+    def __getitem__(self, idx):
+        return self.processed_texts[idx]
+
+
+def batched_collate_fn(batch):
+    x = []
+    x_lengths = []
+
+    for b in batch:
+        x.append(b["x"].squeeze(0))
+        x_lengths.append(b["x_lengths"])
+
+    x = torch.nn.utils.rnn.pad_sequence(x, batch_first=True)
+    x_lengths = torch.concat(x_lengths, dim=0)
+    return {"x": x, "x_lengths": x_lengths}
+
+
+def batched_synthesis(args, device, model, vocoder, denoiser, texts, spk):
+    total_rtf = []
+    total_rtf_w = []
+    processed_text = [process_text(i, text, "cpu") for i, text in enumerate(texts)]
+    dataloader = torch.utils.data.DataLoader(
+        BatchedSynthesisDataset(processed_text),
+        batch_size=args.batch_size,
+        collate_fn=batched_collate_fn,
+        num_workers=8,
+    )
+    for i, batch in enumerate(dataloader):
+        i = i + 1
+        start_t = dt.datetime.now()
+        b = batch["x"].shape[0]
+        output = model.synthesise(
+            batch["x"].to(device),
+            batch["x_lengths"].to(device),
+            n_timesteps=args.steps,
+            temperature=args.temperature,
+            spks=spk.expand(b) if spk is not None else spk,
+            length_scale=args.speaking_rate,
+        )
+
+        output["waveform"] = to_waveform(output["mel"], vocoder, denoiser)
+        t = (dt.datetime.now() - start_t).total_seconds()
+        rtf_w = t * 22050 / (output["waveform"].shape[-1])
+        print(f"[🍵-Batch: {i}] Matcha-TTS RTF: {output['rtf']:.4f}")
+        print(f"[🍵-Batch: {i}] Matcha-TTS + VOCODER RTF: {rtf_w:.4f}")
+        total_rtf.append(output["rtf"])
+        total_rtf_w.append(rtf_w)
+        for j in range(output["mel"].shape[0]):
+            base_name = f"utterance_{j:03d}_speaker_{args.spk:03d}" if args.spk is not None else f"utterance_{j:03d}"
+            length = output["mel_lengths"][j]
+            new_dict = {"mel": output["mel"][j][:, :length], "waveform": output["waveform"][j][: length * 256]}
+            location = save_to_folder(base_name, new_dict, args.output_folder)
+            print(f"[🍵-{j}] Waveform saved: {location}")
+
+    print("".join(["="] * 100))
+    print(f"[🍵] Average Matcha-TTS RTF: {np.mean(total_rtf):.4f} ± {np.std(total_rtf)}")
+    print(f"[🍵] Average Matcha-TTS + VOCODER RTF: {np.mean(total_rtf_w):.4f} ± {np.std(total_rtf_w)}")
+    print("[🍵] Enjoy the freshly whisked 🍵 Matcha-TTS!")
+
+
+def unbatched_synthesis(args, device, model, vocoder, denoiser, texts, spk):
+    total_rtf = []
+    total_rtf_w = []
+    for i, text in enumerate(texts):
+        i = i + 1
+        base_name = f"utterance_{i:03d}_speaker_{args.spk:03d}" if args.spk is not None else f"utterance_{i:03d}"
+
+        print("".join(["="] * 100))
+        text = text.strip()
+        text_processed = process_text(i, text, device)
+
+        print(f"[🍵] Whisking Matcha-T(ea)TS for: {i}")
+        start_t = dt.datetime.now()
+        output = model.synthesise(
+            text_processed["x"],
+            text_processed["x_lengths"],
+            n_timesteps=args.steps,
+            temperature=args.temperature,
+            spks=spk,
+            length_scale=args.speaking_rate,
+        )
+        output["waveform"] = to_waveform(output["mel"], vocoder, denoiser)
+        # RTF with HiFiGAN
+        t = (dt.datetime.now() - start_t).total_seconds()
+        rtf_w = t * 22050 / (output["waveform"].shape[-1])
+        print(f"[🍵-{i}] Matcha-TTS RTF: {output['rtf']:.4f}")
+        print(f"[🍵-{i}] Matcha-TTS + VOCODER RTF: {rtf_w:.4f}")
+        total_rtf.append(output["rtf"])
+        total_rtf_w.append(rtf_w)
+
+        location = save_to_folder(base_name, output, args.output_folder)
+        print(f"[+] Waveform saved: {location}")
+
+    print("".join(["="] * 100))
+    print(f"[🍵] Average Matcha-TTS RTF: {np.mean(total_rtf):.4f} ± {np.std(total_rtf)}")
+    print(f"[🍵] Average Matcha-TTS + VOCODER RTF: {np.mean(total_rtf_w):.4f} ± {np.std(total_rtf_w)}")
+    print("[🍵] Enjoy the freshly whisked 🍵 Matcha-TTS!")
+
+
+def print_config(args):
+    print("[!] Configurations: ")
+    print(f"\t- Model: {args.model}")
+    print(f"\t- Vocoder: {args.vocoder}")
+    print(f"\t- Temperature: {args.temperature}")
+    print(f"\t- Speaking rate: {args.speaking_rate}")
+    print(f"\t- Number of ODE steps: {args.steps}")
+    print(f"\t- Speaker: {args.spk}")
+
+
+def get_device(args):
+    if torch.cuda.is_available() and not args.cpu:
+        print("[+] GPU Available! Using GPU")
+        device = torch.device("cuda")
+    else:
+        print("[-] GPU not available or forced CPU run! Using CPU")
+        device = torch.device("cpu")
+    return device
+
+
+if __name__ == "__main__":
+    cli()

xinference/thirdparty/matcha/data/text_mel_datamodule.py

@@ -0,0 +1,274 @@
+import random
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import numpy as np
+import torch
+import torchaudio as ta
+from lightning import LightningDataModule
+from torch.utils.data.dataloader import DataLoader
+
+from matcha.text import text_to_sequence
+from matcha.utils.audio import mel_spectrogram
+from matcha.utils.model import fix_len_compatibility, normalize
+from matcha.utils.utils import intersperse
+
+
+def parse_filelist(filelist_path, split_char="|"):
+    with open(filelist_path, encoding="utf-8") as f:
+        filepaths_and_text = [line.strip().split(split_char) for line in f]
+    return filepaths_and_text
+
+
+class TextMelDataModule(LightningDataModule):
+    def __init__(  # pylint: disable=unused-argument
+        self,
+        name,
+        train_filelist_path,
+        valid_filelist_path,
+        batch_size,
+        num_workers,
+        pin_memory,
+        cleaners,
+        add_blank,
+        n_spks,
+        n_fft,
+        n_feats,
+        sample_rate,
+        hop_length,
+        win_length,
+        f_min,
+        f_max,
+        data_statistics,
+        seed,
+        load_durations,
+    ):
+        super().__init__()
+
+        # this line allows to access init params with 'self.hparams' attribute
+        # also ensures init params will be stored in ckpt
+        self.save_hyperparameters(logger=False)
+
+    def setup(self, stage: Optional[str] = None):  # pylint: disable=unused-argument
+        """Load data. Set variables: `self.data_train`, `self.data_val`, `self.data_test`.
+
+        This method is called by lightning with both `trainer.fit()` and `trainer.test()`, so be
+        careful not to execute things like random split twice!
+        """
+        # load and split datasets only if not loaded already
+
+        self.trainset = TextMelDataset(  # pylint: disable=attribute-defined-outside-init
+            self.hparams.train_filelist_path,
+            self.hparams.n_spks,
+            self.hparams.cleaners,
+            self.hparams.add_blank,
+            self.hparams.n_fft,
+            self.hparams.n_feats,
+            self.hparams.sample_rate,
+            self.hparams.hop_length,
+            self.hparams.win_length,
+            self.hparams.f_min,
+            self.hparams.f_max,
+            self.hparams.data_statistics,
+            self.hparams.seed,
+            self.hparams.load_durations,
+        )
+        self.validset = TextMelDataset(  # pylint: disable=attribute-defined-outside-init
+            self.hparams.valid_filelist_path,
+            self.hparams.n_spks,
+            self.hparams.cleaners,
+            self.hparams.add_blank,
+            self.hparams.n_fft,
+            self.hparams.n_feats,
+            self.hparams.sample_rate,
+            self.hparams.hop_length,
+            self.hparams.win_length,
+            self.hparams.f_min,
+            self.hparams.f_max,
+            self.hparams.data_statistics,
+            self.hparams.seed,
+            self.hparams.load_durations,
+        )
+
+    def train_dataloader(self):
+        return DataLoader(
+            dataset=self.trainset,
+            batch_size=self.hparams.batch_size,
+            num_workers=self.hparams.num_workers,
+            pin_memory=self.hparams.pin_memory,
+            shuffle=True,
+            collate_fn=TextMelBatchCollate(self.hparams.n_spks),
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            dataset=self.validset,
+            batch_size=self.hparams.batch_size,
+            num_workers=self.hparams.num_workers,
+            pin_memory=self.hparams.pin_memory,
+            shuffle=False,
+            collate_fn=TextMelBatchCollate(self.hparams.n_spks),
+        )
+
+    def teardown(self, stage: Optional[str] = None):
+        """Clean up after fit or test."""
+        pass  # pylint: disable=unnecessary-pass
+
+    def state_dict(self):
+        """Extra things to save to checkpoint."""
+        return {}
+
+    def load_state_dict(self, state_dict: Dict[str, Any]):
+        """Things to do when loading checkpoint."""
+        pass  # pylint: disable=unnecessary-pass
+
+
+class TextMelDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        filelist_path,
+        n_spks,
+        cleaners,
+        add_blank=True,
+        n_fft=1024,
+        n_mels=80,
+        sample_rate=22050,
+        hop_length=256,
+        win_length=1024,
+        f_min=0.0,
+        f_max=8000,
+        data_parameters=None,
+        seed=None,
+        load_durations=False,
+    ):
+        self.filepaths_and_text = parse_filelist(filelist_path)
+        self.n_spks = n_spks
+        self.cleaners = cleaners
+        self.add_blank = add_blank
+        self.n_fft = n_fft
+        self.n_mels = n_mels
+        self.sample_rate = sample_rate
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.f_min = f_min
+        self.f_max = f_max
+        self.load_durations = load_durations
+
+        if data_parameters is not None:
+            self.data_parameters = data_parameters
+        else:
+            self.data_parameters = {"mel_mean": 0, "mel_std": 1}
+        random.seed(seed)
+        random.shuffle(self.filepaths_and_text)
+
+    def get_datapoint(self, filepath_and_text):
+        if self.n_spks > 1:
+            filepath, spk, text = (
+                filepath_and_text[0],
+                int(filepath_and_text[1]),
+                filepath_and_text[2],
+            )
+        else:
+            filepath, text = filepath_and_text[0], filepath_and_text[1]
+            spk = None
+
+        text, cleaned_text = self.get_text(text, add_blank=self.add_blank)
+        mel = self.get_mel(filepath)
+
+        durations = self.get_durations(filepath, text) if self.load_durations else None
+
+        return {"x": text, "y": mel, "spk": spk, "filepath": filepath, "x_text": cleaned_text, "durations": durations}
+
+    def get_durations(self, filepath, text):
+        filepath = Path(filepath)
+        data_dir, name = filepath.parent.parent, filepath.stem
+
+        try:
+            dur_loc = data_dir / "durations" / f"{name}.npy"
+            durs = torch.from_numpy(np.load(dur_loc).astype(int))
+
+        except FileNotFoundError as e:
+            raise FileNotFoundError(
+                f"Tried loading the durations but durations didn't exist at {dur_loc}, make sure you've generate the durations first using: python matcha/utils/get_durations_from_trained_model.py \n"
+            ) from e
+
+        assert len(durs) == len(text), f"Length of durations {len(durs)} and text {len(text)} do not match"
+
+        return durs
+
+    def get_mel(self, filepath):
+        audio, sr = ta.load(filepath)
+        assert sr == self.sample_rate
+        mel = mel_spectrogram(
+            audio,
+            self.n_fft,
+            self.n_mels,
+            self.sample_rate,
+            self.hop_length,
+            self.win_length,
+            self.f_min,
+            self.f_max,
+            center=False,
+        ).squeeze()
+        mel = normalize(mel, self.data_parameters["mel_mean"], self.data_parameters["mel_std"])
+        return mel
+
+    def get_text(self, text, add_blank=True):
+        text_norm, cleaned_text = text_to_sequence(text, self.cleaners)
+        if self.add_blank:
+            text_norm = intersperse(text_norm, 0)
+        text_norm = torch.IntTensor(text_norm)
+        return text_norm, cleaned_text
+
+    def __getitem__(self, index):
+        datapoint = self.get_datapoint(self.filepaths_and_text[index])
+        return datapoint
+
+    def __len__(self):
+        return len(self.filepaths_and_text)
+
+
+class TextMelBatchCollate:
+    def __init__(self, n_spks):
+        self.n_spks = n_spks
+
+    def __call__(self, batch):
+        B = len(batch)
+        y_max_length = max([item["y"].shape[-1] for item in batch])
+        y_max_length = fix_len_compatibility(y_max_length)
+        x_max_length = max([item["x"].shape[-1] for item in batch])
+        n_feats = batch[0]["y"].shape[-2]
+
+        y = torch.zeros((B, n_feats, y_max_length), dtype=torch.float32)
+        x = torch.zeros((B, x_max_length), dtype=torch.long)
+        durations = torch.zeros((B, x_max_length), dtype=torch.long)
+
+        y_lengths, x_lengths = [], []
+        spks = []
+        filepaths, x_texts = [], []
+        for i, item in enumerate(batch):
+            y_, x_ = item["y"], item["x"]
+            y_lengths.append(y_.shape[-1])
+            x_lengths.append(x_.shape[-1])
+            y[i, :, : y_.shape[-1]] = y_
+            x[i, : x_.shape[-1]] = x_
+            spks.append(item["spk"])
+            filepaths.append(item["filepath"])
+            x_texts.append(item["x_text"])
+            if item["durations"] is not None:
+                durations[i, : item["durations"].shape[-1]] = item["durations"]
+
+        y_lengths = torch.tensor(y_lengths, dtype=torch.long)
+        x_lengths = torch.tensor(x_lengths, dtype=torch.long)
+        spks = torch.tensor(spks, dtype=torch.long) if self.n_spks > 1 else None
+
+        return {
+            "x": x,
+            "x_lengths": x_lengths,
+            "y": y,
+            "y_lengths": y_lengths,
+            "spks": spks,
+            "filepaths": filepaths,
+            "x_texts": x_texts,
+            "durations": durations if not torch.eq(durations, 0).all() else None,
+        }