inference.py

# *****************************************************************************
#  Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#      * Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.
#      * Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#      * Neither the name of the NVIDIA CORPORATION nor the
#        names of its contributors may be used to endorse or promote products
#        derived from this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# *****************************************************************************

import argparse
import os
import time
import sys
import warnings

import torch
import torch.distributions as D
import torch.nn.functional as F
import numpy as np
from scipy.stats import norm
from scipy.io.wavfile import write
from torch.nn.utils.rnn import pad_sequence
from tqdm import tqdm

import dllogger as DLLogger
import models
from dllogger import StdOutBackend, JSONStreamBackend, Verbosity

from common.tb_dllogger import (init_inference_metadata, stdout_metric_format,
                                unique_log_fpath)
from common.text.symbols import get_pad_idx
from common.text.text_processing import TextProcessor, PhoneProcessor, UnitProcessor
from common.utils import load_speaker_lang_ids
from hifigan.denoiser import Denoiser
from pitch_transform import pitch_transform_custom


def parse_args(parser):
    """
    Parse commandline arguments.
    """
    parser.add_argument('-i', '--input', type=str, required=True,
                        help='Full path to the input text (phrases separated by newlines)')
    parser.add_argument('-o', '--output', default=None,
                        help='Output folder to save audio (file per phrase)')
    parser.add_argument('-d', '--dataset-path', type=str, default='',
                        help='Path to dataset (for loading extra data fields)')
    parser.add_argument('--log-file', type=str, default=None,
                        help='Path to a DLLogger log file')
    parser.add_argument('--cuda', action='store_true',
                        help='Run inference on a GPU using CUDA')
    parser.add_argument('--cudnn-benchmark', action='store_true',
                        help='Enable cudnn benchmark mode')
    parser.add_argument('--fastpitch', type=str, default='',
                        help='Path to FastPitch checkpoint file (skip to use ground truth mels)')
    parser.add_argument('--hifigan', type=str, default='',
                        help='Path to HiFi-GAN generator checkpoint (skip to only generate mels)')
    parser.add_argument('--hifigan-config', type=str, default='',
                        help='Path to HiFi-GAN config file')
    parser.add_argument('-sr', '--sampling-rate', default=22050, type=int,
                        help='Sampling rate')
    parser.add_argument('--stft-hop-length', type=int, default=256,
                        help='STFT hop length for estimating audio length from mel size')
    parser.add_argument('--denoising-strength', default=0.0, type=float,
                        help='Capture and subtract HiFi-GAN model bias to enhance audio')
    parser.add_argument('--save-mels', action='store_true',
                        help='Save generated mel spectrograms from FastPitch')
    parser.add_argument('--amp', action='store_true',
                        help='Inference with AMP')
    parser.add_argument('-bs', '--batch-size', type=int, default=64)
    parser.add_argument('--warmup-steps', type=int, default=0,
                        help='Warmup iterations before measuring performance')
    parser.add_argument('--repeats', type=int, default=1,
                        help='Repeat inference for benchmarking')
    parser.add_argument('--ema', action='store_true',
                        help='Use EMA averaged model (if saved in checkpoints)')

    transform = parser.add_argument_group('transform')
    transform.add_argument('--fade-out', type=int, default=10,
                           help='Number of fadeout frames at the end')
    transform.add_argument('--pace', type=float, default=1.0,
                           help='Adjust the pace of speech')
    transform.add_argument('--pitch-transform-flatten', action='store_true',
                           help='Flatten the pitch')
    transform.add_argument('--pitch-transform-invert', action='store_true',
                           help='Invert the pitch wrt mean value')
    transform.add_argument('--pitch-transform-amplify', type=float, default=1.0,
                           help='Amplify pitch variability, typical values are in the range (1.0, 3.0).')
    transform.add_argument('--pitch-transform-shift', type=float, default=0.0,
                           help='Raise/lower the pitch by <hz>')
    transform.add_argument('--pitch-transform-custom', action='store_true',
                           help='Apply the transform from pitch_transform.py')

    text_processing = parser.add_argument_group('Text processing parameters')
    text_processing.add_argument('--input-type', type=str, default='char',
                                 choices=['char', 'phone', 'pf', 'unit'],
                                 help='Input symbols used, either char (text), phone, '
                                 'pf (phonological feature vectors) or unit (quantized '
                                 'acoustic representation IDs)')
    text_processing.add_argument('--symbol-set', type=str, default='english_basic',
                                 help='Define symbol set for input sequences. For '
                                 'quantized unit inputs, pass the size of the vocabulary.')
    text_processing.add_argument('--text-cleaners', nargs='*', default=[], type=str,
                                 help='Type of text cleaners for input text.')

    cond = parser.add_argument_group('conditioning on additional attributes')
    cond.add_argument('--speaker-ids', type=str, default=None,
                      help='Speaker ID mapping.')
    # TODO: check, this should work for loading embeddings too
    cond.add_argument('--speaker', type=str, default=None,
                        help='Speaker ID for a multi-speaker model')
    cond.add_argument('--lang-ids', type=str, default=None,
                      help='Language ID mapping.')
    cond.add_argument('--language', type=str, default=None,
                        help='Language ID for a multi-lingual model')

    return parser


def load_model_from_ckpt(checkpoint_path, ema, model, device):
    checkpoint_data = torch.load(checkpoint_path, map_location=device)
    status = ''

    if 'state_dict' in checkpoint_data:
        # FastPitch checkpoint
        sd = checkpoint_data['state_dict']
        if ema and 'ema_state_dict' in checkpoint_data:
            sd = checkpoint_data['ema_state_dict']
            status += ' (EMA)'
        elif ema and not 'ema_state_dict' in checkpoint_data:
            print(f'WARNING: EMA weights missing for {checkpoint_data}')

        if any(key.startswith('module.') for key in sd):
            sd = {k.replace('module.', ''): v for k,v in sd.items()}
        status += ' ' + str(model.load_state_dict(sd))
    elif 'generator' in checkpoint_data:
        # HiFi-GAN checkpoint
        model.load_state_dict(checkpoint_data['generator'])
    else:
        model = checkpoint_data['model']

    print(f'Loaded {checkpoint_path}{status}')
    return model


def load_and_setup_model(model_name, parser, checkpoint, amp, device,
                         unk_args=[], forward_is_infer=False, ema=True):
                         
    model_parser = models.parse_model_args(model_name, parser, add_help=False)
    model_args, model_unk_args = model_parser.parse_known_args()
    unk_args[:] = list(set(unk_args) & set(model_unk_args))

    model_config = models.get_model_config(model_name, model_args)

    model = models.get_model(model_name, model_config, device,
                             forward_is_infer=forward_is_infer)

    if checkpoint is not None:
        model = load_model_from_ckpt(checkpoint, ema, model, device)

    if model_name == "HiFi-GAN":
        model.remove_weight_norm()

    if amp:
        model.half()
    model.eval()
    return model.to(device)


def load_fields(fpath):
    lines = [l.strip() for l in open(fpath, encoding='utf-8')]
    columns = lines[0].split('|')
    fields = list(zip(*[t.split('|') for t in lines[1:]]))
    return {c:f for c, f in zip(columns, fields)}


def prepare_input_sequence(fields, device, input_type, symbol_set, text_cleaners,
                           batch_size=128, dataset=None,
                           load_mels=False, load_pitch=False, load_duration=False,
                           load_speaker=False, speaker_ids=None,
                           load_lang=False, lang_ids=None):
    if input_type == 'char':
        tp = TextProcessor(symbol_set, text_cleaners)
    elif input_type == 'unit':
        tp = UnitProcessor(symbol_set, input_type)
    else:
        tp = PhoneProcessor(symbol_set, input_type)

    if input_type == 'pf':
        fields['text'] = [torch.FloatTensor(tp.encode_text(text))
                          for text in tqdm(fields['text'], "Loading texts")]
    else:
        fields['text'] = [torch.LongTensor(tp.encode_text(text))
                          for text in tqdm(fields['text'], "Loading texts")]
    order = np.argsort([-t.size(0) for t in fields['text']])

    fields['text'] = [fields['text'][i] for i in order]
    fields['text_lens'] = torch.LongTensor([t.size(0) for t in fields['text']])

    if load_mels:
        fields['mel'] = [
            torch.load(os.path.join(dataset, fields['mel'][i])).t()
            for i in tqdm(order, "Loading mels")]
        fields['mel_lens'] = torch.LongTensor([t.size(0) for t in fields['mel']])
    if load_pitch:
        fields['pitch'] = [
            torch.load(os.path.join(dataset, fields['pitch'][i])).float()
            for i in tqdm(order, "Loading pitches")]
        fields['pitch_lens'] = torch.LongTensor([t.size(0) for t in fields['pitch']])
    if load_duration:
        fields['duration'] = [
            torch.load(os.path.join(dataset, fields['duration'][i]))
            for i in tqdm(order, "Loading durations")]
    if load_speaker:
        if speaker_ids is not None:
            fields['speaker'] = torch.LongTensor([
                speaker_ids[fields['speaker'][i]] for i in order])
        else:
            fields['speaker'] = torch.cat([
                torch.load(os.path.join(dataset, fields['speaker'][i]))
                for i in tqdm(order, "Loading speaker embeddings")])
    if load_lang:
        if lang_ids is not None:
            fields['language'] = torch.LongTensor([
                lang_ids[fields['language'][i]] for i in order])
        else:
            fields['language'] = torch.cat([
                torch.load(os.path.join(dataset, fields['language'][i]))
                for i in tqdm(order, "Loading language embeddings")])
    if 'output' in fields:
        fields['output'] = [fields['output'][i] for i in order]
    if 'mel_output' in fields:
        fields['mel_output'] = [fields['mel_output'][i] for i in order]

    # cut into batches & pad
    batches = []
    for b in range(0, len(order), batch_size):
        batch = {f: values[b:b+batch_size] for f, values in fields.items()}
        for f in batch:
            if f == 'text':
                batch[f] = pad_sequence(batch[f], batch_first=True,
                                        padding_value=get_pad_idx(symbol_set, input_type))
            elif f == 'mel' and load_mels:
                batch[f] = pad_sequence(batch[f], batch_first=True).permute(0, 2, 1)
            elif f == 'pitch' and load_pitch:
                batch[f] = pad_sequence(batch[f], batch_first=True)
            elif f == 'duration' and load_duration:
                if 'mel' in batch:
                    # we're in mas mode: use gt mel and text lens to drive
                    # masked attention and ensure final durations are right
                    # (individual segment durations are still predicted --
                    # we hope they line up enough for vocoder fine-tuning
                    # to make sense)
                    attn_prior_dims = [i.shape for i in batch[f]]
                    mel_lens = [i[0] for i in attn_prior_dims]
                    text_lens = [i[1] for i in attn_prior_dims]
                    dur_padded = torch.zeros(
                        len(mel_lens), max(mel_lens), max(text_lens)).to(device)
                    for i, attn_prior in enumerate(batch[f]):
                        dur_padded[i, :attn_prior.size(0), :attn_prior.size(1)] = attn_prior
                    mel_lens = torch.as_tensor(mel_lens, dtype=torch.int).to(device)
                    text_lens = torch.as_tensor(text_lens, dtype=torch.int).to(device)
                    batch[f] = (dur_padded, batch['mel'].to(device), mel_lens, text_lens)
                else:
                    batch[f] = pad_sequence(batch[f], batch_first=True)

            if type(batch[f]) is torch.Tensor:
                batch[f] = batch[f].to(device)
        batches.append(batch)
    return batches


def build_pitch_transformation(args):
    if args.pitch_transform_custom:
        def custom_(pitch, pitch_lens, mean, std):
            return (pitch_transform_custom(pitch * std + mean, pitch_lens)
                    - mean) / std
        return custom_

    fun = 'pitch'
    if args.pitch_transform_flatten:
        fun = f'({fun}) * 0.0'
    if args.pitch_transform_invert:
        fun = f'({fun}) * -1.0'
    if args.pitch_transform_amplify:
        ampl = args.pitch_transform_amplify
        fun = f'({fun}) * {ampl}'
    if args.pitch_transform_shift != 0.0:
        hz = args.pitch_transform_shift
        fun = f'({fun}) + {hz} / std'
    return eval(f'lambda pitch, pitch_lens, mean, std: {fun}')


class MeasureTime(list):
    def __init__(self, *args, cuda=True, **kwargs):
        super(MeasureTime, self).__init__(*args, **kwargs)
        self.cuda = cuda

    def __enter__(self):
        if self.cuda:
            torch.cuda.synchronize()
        self.t0 = time.perf_counter()

    def __exit__(self, exc_type, exc_value, exc_traceback):
        if self.cuda:
            torch.cuda.synchronize()
        self.append(time.perf_counter() - self.t0)

    def __add__(self, other):
        assert len(self) == len(other)
        return MeasureTime((sum(ab) for ab in zip(self, other)), cuda=cuda)


def main():
    """
    Launches text to speech (inference).
    Inference is executed on a single GPU.
    """
    parser = argparse.ArgumentParser(description='PyTorch FastPitch Inference',
                                     allow_abbrev=False)
    parser = parse_args(parser)
    args, unk_args = parser.parse_known_args()

    torch.backends.cudnn.benchmark = args.cudnn_benchmark

    if args.output is not None:
        os.makedirs(args.output, exist_ok=True)

    log_fpath = args.log_file or os.path.join(args.output, 'nvlog_infer.json')
    log_fpath = unique_log_fpath(log_fpath)
    DLLogger.init(backends=[JSONStreamBackend(Verbosity.DEFAULT, log_fpath),
                            StdOutBackend(Verbosity.VERBOSE,
                                          metric_format=stdout_metric_format)])
    init_inference_metadata()
    [DLLogger.log("PARAMETER", {k: v}) for k, v in vars(args).items()]

    device = torch.device('cuda' if args.cuda else 'cpu')

    generator = None
    vocoder = None
    denoiser = None

    if args.fastpitch:
        generator = load_and_setup_model(
            'FastPitch', parser, args.fastpitch, args.amp, device,
            unk_args=unk_args, forward_is_infer=True, ema=args.ema)

    if args.hifigan:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            vocoder = load_and_setup_model(
                'HiFi-GAN', parser, args.hifigan, args.amp, device,
                unk_args=unk_args, forward_is_infer=True, ema=args.ema)
        if args.denoising_strength > 0.0:
            denoiser = Denoiser(
                vocoder, sr=args.sampling_rate, hop_length=args.stft_hop_length).to(device)

    if len(unk_args) > 0:
        raise ValueError(f'Invalid options {unk_args}')

    speaker_ids = load_speaker_lang_ids(args.speaker_ids)
    lang_ids = load_speaker_lang_ids(args.lang_ids)

    fields = load_fields(args.input)
    # override from command line
    if args.speaker is not None:
        fields['speaker'] = list(args.speaker for _ in fields['text'])
    if args.language is not None:
        fields['language'] = list(args.language for _ in fields['text'])

    batches = prepare_input_sequence(
        fields, device, args.input_type, args.symbol_set, args.text_cleaners,
        args.batch_size, args.dataset_path, load_mels=(generator is None or 'mel' in fields),
        load_pitch=('pitch' in fields), load_duration=('duration' in fields),
        load_speaker=('speaker' in fields), speaker_ids=speaker_ids,
        load_lang=('language' in fields), lang_ids=lang_ids)

    # Use real data rather than synthetic - FastPitch predicts len
    if args.warmup_steps:
        for _ in tqdm(range(args.warmup_steps), 'Warmup'):
            with torch.no_grad():
                if generator is not None:
                    b = batches[0]
                    mel, *_ = generator(b['text'])

                if vocoder is not None:
                    audios = vocoder(mel)

    gen_measures = MeasureTime(cuda=args.cuda)
    vocoder_measures = MeasureTime(cuda=args.cuda)

    gen_kw = {'pace': args.pace,
              'speaker': args.speaker,
              'language': args.language,
              'pitch_transform': build_pitch_transformation(args)}

    all_utterances = 0
    all_samples = 0
    all_letters = 0
    all_frames = 0

    reps = args.repeats
    log_enabled = reps == 1
    log = lambda s, d: DLLogger.log(step=s, data=d) if log_enabled else None

    for rep in (tqdm(range(reps), 'Inference') if reps > 1 else range(reps)):
        for n, b in enumerate(tqdm(batches, "Synthesizing utterances")):
            if generator is None:
                mel, mel_lens = b['mel'], b['mel_lens']
            else:
                gen_kw['dur_tgt'] = b['duration'] if 'duration' in b else None
                gen_kw['pitch_tgt'] = b['pitch'] if 'pitch' in b else None
                gen_kw['speaker'] = b['speaker'] if 'speaker' in b else args.speaker
                gen_kw['language'] = b['language'] if 'language' in b else args.language
                with torch.no_grad(), gen_measures:
                    mel, mel_lens, *_ = generator(b['text'], **gen_kw)

                    if generator.tvcgmm_k:
                        min_var = 1.0e-3
                        n_mel = generator.n_mel_channels
                        param_predictions = mel.transpose(1, 2).reshape(
                             args.batch_size, -1, n_mel, generator.tvcgmm_k, 10)

                        pis = F.softmax(param_predictions[..., 0], dim=-1)
                        mus = param_predictions[..., 1:4]
                        scale_trils = torch.diag_embed(
                            F.softplus(param_predictions[..., 4:7]) + min_var, offset=0)
                        scale_trils += torch.diag_embed(param_predictions[..., 7:9], offset=-1)
                        scale_trils += torch.diag_embed(param_predictions[..., 9:10], offset=-2)
                        sigmas = scale_trils @ scale_trils.transpose(-1, -2)
                        # last frequency bin is erroneous in conditional sampling due to 
                        # missing training targets
                        sigmas[:, -1] = sigmas[:, -2]

                        mix = D.Categorical(pis)
                        comp = D.MultivariateNormal(mus, scale_tril=scale_trils)
                        mixture = D.MixtureSameFamily(mix, comp)

                        mel_pred = mixture.sample().transpose(2, 3)
                        mel_pred = mel_pred.reshape(args.batch_size, -1, 3 * n_mel).transpose(1, 2)
                        mel_pred[:, :n_mel, 1:] += mel_pred[:, n_mel:2 * n_mel, :-1]
                        mel_pred[:, 1:n_mel, :] += mel_pred[:, 2 * n_mel:-1, :]
                        mel_pred[:, 1:n_mel, 1:] /= 3
                        mel_pred[:, 0, 1:] /= 2
                        mel_pred[:, 1:, 0] /= 2
                        mel_pred = mel_pred[:, :n_mel]
                        mel = mel_pred

                gen_infer_perf = mel.size(0) * mel.size(2) / gen_measures[-1]
                all_letters += b['text_lens'].sum().item()
                all_frames += mel.size(0) * mel.size(2)

                if args.save_mels and args.output is not None and reps == 1:
                    for i, _mel in enumerate(mel):
                        _mel = _mel[:, :mel_lens[i]]
                        if 'mel_output' in b:
                            fname = b['mel_output'][i]
                        elif 'output' in b:
                            fname = os.path.splitext(b['output'][i])[0] + '.pt'
                        else:
                            fname = f'mel_{all_utterances + i}.pt'
                        mel_path = os.path.join(args.output, fname)
                        torch.save(_mel.clone().cpu(), mel_path)

            if vocoder is not None:
                with torch.no_grad(), vocoder_measures:
                    audios = vocoder(mel)
                    if denoiser is not None:
                        audios = denoiser(audios.squeeze(1), args.denoising_strength)
                    audios = audios.squeeze(1)

                all_samples += sum(audio.size(0) for audio in audios)
                vocoder_infer_perf = (
                    audios.size(0) * audios.size(1) / vocoder_measures[-1])

                if args.output is not None and reps == 1:
                    for i, audio in enumerate(audios):
                        audio = audio[:mel_lens[i].item() * args.stft_hop_length]

                        if args.fade_out:
                            fade_len = args.fade_out * args.stft_hop_length
                            fade_w = torch.linspace(1.0, 0.0, fade_len)
                            audio[-fade_len:] *= fade_w.to(audio.device)

                        audio = audio / torch.max(torch.abs(audio))
                        fname = b['output'][i] if 'output' in b else f'audio_{all_utterances + i}.wav'
                        audio_path = os.path.join(args.output, fname)
                        write(audio_path, args.sampling_rate, audio.cpu().numpy())

                all_utterances += len(audios)

    log_enabled = True
    if generator is not None:
        gm = np.sort(np.asarray(gen_measures))
        rtf = all_samples / (all_utterances * gm.mean() * args.sampling_rate)
        log((), {"avg_fastpitch_letters/s": all_letters / gm.sum()})
        log((), {"avg_fastpitch_frames/s": all_frames / gm.sum()})
        log((), {"avg_fastpitch_latency": gm.mean()})
        log((), {"avg_fastpitch_RTF": rtf})
        log((), {"90%_fastpitch_latency": gm.mean() + norm.ppf((1.0 + 0.90) / 2) * gm.std()})
        log((), {"95%_fastpitch_latency": gm.mean() + norm.ppf((1.0 + 0.95) / 2) * gm.std()})
        log((), {"99%_fastpitch_latency": gm.mean() + norm.ppf((1.0 + 0.99) / 2) * gm.std()})
    if vocoder is not None:
        wm = np.sort(np.asarray(vocoder_measures))
        rtf = all_samples / (all_utterances * wm.mean() * args.sampling_rate)
        log((), {"avg_vocoder_samples/s": all_samples / wm.sum()})
        log((), {"avg_vocoder_latency": wm.mean()})
        log((), {"avg_vocoder_RTF": rtf})
        log((), {"90%_vocoder_latency": wm.mean() + norm.ppf((1.0 + 0.90) / 2) * wm.std()})
        log((), {"95%_vocoder_latency": wm.mean() + norm.ppf((1.0 + 0.95) / 2) * wm.std()})
        log((), {"99%_vocoder_latency": wm.mean() + norm.ppf((1.0 + 0.99) / 2) * wm.std()})
    if generator is not None and vocoder is not None:
        m = gm + wm
        rtf = all_samples / (all_utterances * m.mean() * args.sampling_rate)
        log((), {"avg_samples/s": all_samples / m.sum()})
        log((), {"avg_letters/s": all_letters / m.sum()})
        log((), {"avg_latency": m.mean()})
        log((), {"avg_RTF": rtf})
        log((), {"90%_latency": m.mean() + norm.ppf((1.0 + 0.90) / 2) * m.std()})
        log((), {"95%_latency": m.mean() + norm.ppf((1.0 + 0.95) / 2) * m.std()})
        log((), {"99%_latency": m.mean() + norm.ppf((1.0 + 0.99) / 2) * m.std()})
    DLLogger.flush()


if __name__ == '__main__':
    main()