scripts/dataset_processing/tts/resynthesize_dataset.py

# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
This script is a helper for resynthesizing TTS dataset using a pretrained text-to-spectrogram model.
Goal of resynthesis (as opposed to text-to-speech) is to use the largest amount of ground-truth features from existing speech data.
For example, for resynthesis we want to have the same pitch and durations instead of ones predicted by the model.
The results are to be used for some other task: vocoder finetuning, spectrogram enhancer training, etc.

Let's say we have the following toy dataset:
/dataset/manifest.json
/dataset/1/foo.wav
/dataset/2/bar.wav
/dataset/sup_data/pitch/1_foo.pt
/dataset/sup_data/pitch/2_bar.pt

manifest.json has two entries for "/dataset/1/foo.wav" and "/dataset/2/bar.wav"
(sup_data folder contains pitch files precomputed during training a FastPitch model on this dataset.)
(If you lost your sup_data - don't worry, we use TTSDataset class so they would be created on-the-fly)

Our script call is
$ python scripts/dataset_processing/tts/resynthesize_dataset.py \
    --model-path ./models/fastpitch/multi_spk/FastPitch--val_loss\=1.4473-epoch\=209.ckpt \
    --input-json-manifest "/dataset/manifest.json" \
    --input-sup-data-path "/dataset/sup_data/" \
    --output-folder "/output/" \
    --device "cuda:0" \
    --batch-size 1 \
    --num-workers 1

Then we get output dataset with following directory structure:
/output/manifest_mel.json
/output/mels/foo.npy
/output/mels/foo_gt.npy
/output/mels/bar.npy
/output/mels/bar_gt.npy

/output/manifest_mel.json has the same entries as /dataset/manifest.json but with new fields for spectrograms.
"mel_filepath" is path to the resynthesized spectrogram .npy, "mel_gt_filepath" is path to ground-truth spectrogram .npy

The output structure is similar to generate_mels.py script for compatibility reasons.
"""

import argparse
import itertools
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, Iterator, List

import numpy as np
import torch
from omegaconf import DictConfig, OmegaConf
from tqdm import tqdm

from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest
from nemo.collections.tts.models import FastPitchModel
from nemo.collections.tts.models.base import SpectrogramGenerator
from nemo.collections.tts.parts.utils.helpers import process_batch, to_device_recursive


def chunks(iterable: Iterable, size: int) -> Iterator[List]:
    # chunks([1, 2, 3, 4, 5], size=2) -> [[1, 2], [3, 4], [5]]
    # assumes iterable does not have any `None`s
    args = [iter(iterable)] * size
    for chunk in itertools.zip_longest(*args, fillvalue=None):
        chunk = list(item for item in chunk if item is not None)
        if chunk:
            yield chunk


def load_model(path: Path, device: torch.device) -> SpectrogramGenerator:
    model = None
    if path.suffix == ".nemo":
        model = SpectrogramGenerator.restore_from(path, map_location=device)
    elif path.suffix == ".ckpt":
        model = SpectrogramGenerator.load_from_checkpoint(path, map_location=device)
    else:
        raise ValueError(f"Unknown checkpoint type {path.suffix} ({path})")

    return model.eval().to(device)


@dataclass
class TTSDatasetResynthesizer:
    """
    Reuses internals of a SpectrogramGenerator to resynthesize dataset using ground truth features.
    Default setup is FastPitch with learned alignment.
    If your use case requires different setup, you can either contribute to this script or subclass this class.
    """

    model: SpectrogramGenerator
    device: torch.device

    @torch.no_grad()
    def resynthesize_batch(self, batch: Dict[str, Any]) -> Dict[str, Any]:
        """
        Resynthesizes a single batch.
        Takes a dict with main data and sup data.
        Outputs a dict with model outputs.
        """
        if not isinstance(self.model, FastPitchModel):
            raise NotImplementedError(
                "This script supports only FastPitch. Please implement resynthesizing routine for your desired model."
            )

        batch = to_device_recursive(batch, self.device)

        mels, mel_lens = self.model.preprocessor(input_signal=batch["audio"], length=batch["audio_lens"])
        outputs_tuple = self.model.forward(
            text=batch["text"],
            durs=None,
            pitch=batch["pitch"],
            speaker=batch.get("speaker"),
            pace=1.0,
            spec=mels,
            attn_prior=batch.get("attn_prior"),
            mel_lens=mel_lens,
            input_lens=batch["text_lens"],
        )
        names = self.model.fastpitch.output_types.keys()
        return {"spec": mels, "mel_lens": mel_lens, **dict(zip(names, outputs_tuple))}

    def resynthesized_batches(self) -> Iterator[Dict[str, Any]]:
        """
        Returns a generator of resynthesized batches.
        Each returned batch is a dict containing main data, sup data, and model output
        """
        self.model.setup_training_data(self.model._cfg["train_ds"])

        for batch_tuple in iter(self.model._train_dl):
            batch = process_batch(batch_tuple, sup_data_types_set=self.model._train_dl.dataset.sup_data_types)
            yield self.resynthesize_batch(batch)


def prepare_paired_mel_spectrograms(
    model_path: Path,
    input_json_manifest: Path,
    input_sup_data_path: Path,
    output_folder: Path,
    device: torch.device,
    batch_size: int,
    num_workers: int,
):
    model = load_model(model_path, device)

    dataset_config_overrides = {
        "dataset": {
            "manifest_filepath": str(input_json_manifest.absolute()),
            "sup_data_path": str(input_sup_data_path.absolute()),
        },
        "dataloader_params": {"batch_size": batch_size, "num_workers": num_workers, "shuffle": False},
    }
    model._cfg.train_ds = OmegaConf.merge(model._cfg.train_ds, DictConfig(dataset_config_overrides))
    resynthesizer = TTSDatasetResynthesizer(model, device)

    input_manifest = read_manifest(input_json_manifest)

    output_manifest = []
    output_json_manifest = output_folder / f"{input_json_manifest.stem}_mel{input_json_manifest.suffix}"
    output_mels_folder = output_folder / "mels"
    output_mels_folder.mkdir(exist_ok=True, parents=True)
    for batch, batch_manifest in tqdm(
        zip(resynthesizer.resynthesized_batches(), chunks(input_manifest, size=batch_size)), desc="Batch #"
    ):
        pred_mels = batch["spect"].cpu()  # key from fastpitch.output_types
        true_mels = batch["spec"].cpu()  # key from code above
        mel_lens = batch["mel_lens"].cpu().flatten()  # key from code above

        for i, (manifest_entry, length) in enumerate(zip(batch_manifest, mel_lens.tolist())):
            print(manifest_entry["audio_filepath"])
            filename = Path(manifest_entry["audio_filepath"]).stem

            # note that lengths match
            pred_mel = pred_mels[i, :, :length].clone().numpy()
            true_mel = true_mels[i, :, :length].clone().numpy()

            pred_mel_path = output_mels_folder / f"{filename}.npy"
            true_mel_path = output_mels_folder / f"{filename}_gt.npy"

            np.save(pred_mel_path, pred_mel)
            np.save(true_mel_path, true_mel)

            new_manifest_entry = {
                **manifest_entry,
                "mel_filepath": str(pred_mel_path),
                "mel_gt_filepath": str(true_mel_path),
            }
            output_manifest.append(new_manifest_entry)

    write_manifest(output_json_manifest, output_manifest, ensure_ascii=False)


def argument_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Resynthesize TTS dataset using a pretrained text-to-spectrogram model",
    )
    parser.add_argument(
        "--model-path", required=True, type=Path, help="Path to a checkpoint (either .nemo or .ckpt)",
    )
    parser.add_argument(
        "--input-json-manifest", required=True, type=Path, help="Path to the input JSON manifest",
    )
    parser.add_argument(
        "--input-sup-data-path", required=True, type=Path, help="sup_data_path for the JSON manifest",
    )
    parser.add_argument(
        "--output-folder",
        required=True,
        type=Path,
        help="Path to the output folder. Will contain updated manifest and mels/ folder with spectrograms in .npy files",
    )
    parser.add_argument("--device", required=True, type=torch.device, help="Device ('cpu', 'cuda:0', ...)")
    parser.add_argument("--batch-size", required=True, type=int, help="Batch size in the DataLoader")
    parser.add_argument("--num-workers", required=True, type=int, help="Num workers in the DataLoader")
    return parser


if __name__ == "__main__":
    arguments = argument_parser().parse_args()
    prepare_paired_mel_spectrograms(**vars(arguments))