fairseq/examples/speech_synthesis/generate_waveform.py

# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.


import logging
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path
import soundfile as sf
import sys
import torch
import torchaudio

from fairseq import checkpoint_utils, options, tasks, utils
from fairseq.logging import progress_bar
from fairseq.tasks.text_to_speech import plot_tts_output
from fairseq.data.audio.text_to_speech_dataset import TextToSpeechDataset


logging.basicConfig()
logging.root.setLevel(logging.INFO)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def make_parser():
    parser = options.get_speech_generation_parser()
    parser.add_argument("--dump-features", action="store_true")
    parser.add_argument("--dump-waveforms", action="store_true")
    parser.add_argument("--dump-attentions", action="store_true")
    parser.add_argument("--dump-eos-probs", action="store_true")
    parser.add_argument("--dump-plots", action="store_true")
    parser.add_argument("--dump-target", action="store_true")
    parser.add_argument("--output-sample-rate", default=22050, type=int)
    parser.add_argument("--teacher-forcing", action="store_true")
    parser.add_argument(
        "--audio-format", type=str, default="wav", choices=["wav", "flac"]
    )
    return parser


def postprocess_results(
        dataset: TextToSpeechDataset, sample, hypos, resample_fn, dump_target
):
    def to_np(x):
        return None if x is None else x.detach().cpu().numpy()

    sample_ids = [dataset.ids[i] for i in sample["id"].tolist()]
    texts = sample["src_texts"]
    attns = [to_np(hypo["attn"]) for hypo in hypos]
    eos_probs = [to_np(hypo.get("eos_prob", None)) for hypo in hypos]
    feat_preds = [to_np(hypo["feature"]) for hypo in hypos]
    wave_preds = [to_np(resample_fn(h["waveform"])) for h in hypos]
    if dump_target:
        feat_targs = [to_np(hypo["targ_feature"]) for hypo in hypos]
        wave_targs = [to_np(resample_fn(h["targ_waveform"])) for h in hypos]
    else:
        feat_targs = [None for _ in hypos]
        wave_targs = [None for _ in hypos]

    return zip(sample_ids, texts, attns, eos_probs, feat_preds, wave_preds,
               feat_targs, wave_targs)


def dump_result(
        is_na_model,
        args,
        vocoder,
        sample_id,
        text,
        attn,
        eos_prob,
        feat_pred,
        wave_pred,
        feat_targ,
        wave_targ,
):
    sample_rate = args.output_sample_rate
    out_root = Path(args.results_path)
    if args.dump_features:
        feat_dir = out_root / "feat"
        feat_dir.mkdir(exist_ok=True, parents=True)
        np.save(feat_dir / f"{sample_id}.npy", feat_pred)
        if args.dump_target:
            feat_tgt_dir = out_root / "feat_tgt"
            feat_tgt_dir.mkdir(exist_ok=True, parents=True)
            np.save(feat_tgt_dir / f"{sample_id}.npy", feat_targ)
    if args.dump_attentions:
        attn_dir = out_root / "attn"
        attn_dir.mkdir(exist_ok=True, parents=True)
        np.save(attn_dir / f"{sample_id}.npy", attn.numpy())
    if args.dump_eos_probs and not is_na_model:
        eos_dir = out_root / "eos"
        eos_dir.mkdir(exist_ok=True, parents=True)
        np.save(eos_dir / f"{sample_id}.npy", eos_prob)

    if args.dump_plots:
        images = [feat_pred.T] if is_na_model else [feat_pred.T, attn]
        names = ["output"] if is_na_model else ["output", "alignment"]
        if feat_targ is not None:
            images = [feat_targ.T] + images
            names = [f"target (idx={sample_id})"] + names
        if is_na_model:
            plot_tts_output(images, names, attn, "alignment", suptitle=text)
        else:
            plot_tts_output(images, names, eos_prob, "eos prob", suptitle=text)
        plot_dir = out_root / "plot"
        plot_dir.mkdir(exist_ok=True, parents=True)
        plt.savefig(plot_dir / f"{sample_id}.png")
        plt.close()

    if args.dump_waveforms:
        ext = args.audio_format
        if wave_pred is not None:
            wav_dir = out_root / f"{ext}_{sample_rate}hz_{vocoder}"
            wav_dir.mkdir(exist_ok=True, parents=True)
            sf.write(wav_dir / f"{sample_id}.{ext}", wave_pred, sample_rate)
        if args.dump_target and wave_targ is not None:
            wav_tgt_dir = out_root / f"{ext}_{sample_rate}hz_{vocoder}_tgt"
            wav_tgt_dir.mkdir(exist_ok=True, parents=True)
            sf.write(wav_tgt_dir / f"{sample_id}.{ext}", wave_targ, sample_rate)


def main(args):
    assert(args.dump_features or args.dump_waveforms or args.dump_attentions
           or args.dump_eos_probs or args.dump_plots)
    if args.max_tokens is None and args.batch_size is None:
        args.max_tokens = 8000
    logger.info(args)

    use_cuda = torch.cuda.is_available() and not args.cpu
    task = tasks.setup_task(args)
    models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
        [args.path],
        task=task,
    )
    model = models[0].cuda() if use_cuda else models[0]
    # use the original n_frames_per_step
    task.args.n_frames_per_step = saved_cfg.task.n_frames_per_step
    task.load_dataset(args.gen_subset, task_cfg=saved_cfg.task)

    data_cfg = task.data_cfg
    sample_rate = data_cfg.config.get("features", {}).get("sample_rate", 22050)
    resample_fn = {
        False: lambda x: x,
        True: lambda x: torchaudio.sox_effects.apply_effects_tensor(
            x.detach().cpu().unsqueeze(0), sample_rate,
            [['rate', str(args.output_sample_rate)]]
        )[0].squeeze(0)
    }.get(args.output_sample_rate != sample_rate)
    if args.output_sample_rate != sample_rate:
        logger.info(f"resampling to {args.output_sample_rate}Hz")

    generator = task.build_generator([model], args)
    itr = task.get_batch_iterator(
        dataset=task.dataset(args.gen_subset),
        max_tokens=args.max_tokens,
        max_sentences=args.batch_size,
        max_positions=(sys.maxsize, sys.maxsize),
        ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test,
        required_batch_size_multiple=args.required_batch_size_multiple,
        num_shards=args.num_shards,
        shard_id=args.shard_id,
        num_workers=args.num_workers,
        data_buffer_size=args.data_buffer_size,
    ).next_epoch_itr(shuffle=False)

    Path(args.results_path).mkdir(exist_ok=True, parents=True)
    is_na_model = getattr(model, "NON_AUTOREGRESSIVE", False)
    dataset = task.dataset(args.gen_subset)
    vocoder = task.args.vocoder
    with progress_bar.build_progress_bar(args, itr) as t:
        for sample in t:
            sample = utils.move_to_cuda(sample) if use_cuda else sample
            hypos = generator.generate(model, sample, has_targ=args.dump_target)
            for result in postprocess_results(
                    dataset, sample, hypos, resample_fn, args.dump_target
            ):
                dump_result(is_na_model, args, vocoder, *result)


def cli_main():
    parser = make_parser()
    args = options.parse_args_and_arch(parser)
    main(args)


if __name__ == "__main__":
    cli_main()