attribute_prediction_model.py

# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: MIT
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.
import torch
from torch import nn
from common import ConvNorm, Invertible1x1Conv
from common import AffineTransformationLayer, SplineTransformationLayer
from common import ConvLSTMLinear
from transformer import FFTransformer
from autoregressive_flow import AR_Step, AR_Back_Step


def get_attribute_prediction_model(config):
    name = config['name']
    hparams = config['hparams']
    if name == 'dap':
        model = DAP(**hparams)
    elif name == 'bgap':
        model = BGAP(**hparams)
    elif name == 'agap':
        model = AGAP(**hparams)
    else:
        raise Exception("{} model is not supported".format(name))

    return model


class AttributeProcessing():
    def __init__(self, take_log_of_input=False):
        super(AttributeProcessing).__init__()
        self.take_log_of_input = take_log_of_input

    def normalize(self, x):
        if self.take_log_of_input:
            x = torch.log(x + 1)
        return x

    def denormalize(self, x):
        if self.take_log_of_input:
            x = torch.exp(x) - 1
        return x


class BottleneckLayerLayer(nn.Module):
    def __init__(self, in_dim, reduction_factor, norm='weightnorm',
                 non_linearity='relu', kernel_size=3, use_partial_padding=False):
        super(BottleneckLayerLayer, self).__init__()

        self.reduction_factor = reduction_factor
        reduced_dim = int(in_dim / reduction_factor)
        self.out_dim = reduced_dim
        if self.reduction_factor > 1:
            fn = ConvNorm(in_dim, reduced_dim, kernel_size=kernel_size,
                          use_weight_norm=(norm == 'weightnorm'))
            if norm == 'instancenorm':
                fn = nn.Sequential(
                    fn, nn.InstanceNorm1d(reduced_dim, affine=True))

            self.projection_fn = fn
            self.non_linearity = nn.ReLU()
            if non_linearity == 'leakyrelu':
                self.non_linearity= nn.LeakyReLU()

    def forward(self, x):
        if self.reduction_factor > 1:
            x = self.projection_fn(x)
            x = self.non_linearity(x)
        return x


class DAP(nn.Module):
    def __init__(self, n_speaker_dim, bottleneck_hparams, take_log_of_input,
                 arch_hparams, use_transformer=False):
        super(DAP, self).__init__()
        self.attribute_processing = AttributeProcessing(take_log_of_input)
        self.bottleneck_layer = BottleneckLayerLayer(**bottleneck_hparams)

        arch_hparams['in_dim'] = self.bottleneck_layer.out_dim + n_speaker_dim
        if use_transformer:
            self.feat_pred_fn = FFTransformer(**arch_hparams)
        else:
            self.feat_pred_fn = ConvLSTMLinear(**arch_hparams)

    def forward(self, txt_enc, spk_emb, x, lens):
        if x is not None:
            x = self.attribute_processing.normalize(x)

        txt_enc = self.bottleneck_layer(txt_enc)
        spk_emb_expanded = spk_emb[..., None].expand(-1, -1, txt_enc.shape[2])
        context = torch.cat((txt_enc, spk_emb_expanded), 1)

        x_hat = self.feat_pred_fn(context, lens)

        outputs = {'x_hat': x_hat, 'x': x}
        return outputs

    def infer(self, z, txt_enc, spk_emb, lens=None):
        x_hat = self.forward(txt_enc, spk_emb, x=None, lens=lens)['x_hat']
        x_hat = self.attribute_processing.denormalize(x_hat)
        return x_hat


class BGAP(torch.nn.Module):
    def __init__(self, n_in_dim, n_speaker_dim, bottleneck_hparams, n_flows,
                 n_group_size, n_layers, with_dilation,
                 kernel_size, scaling_fn,
                 take_log_of_input=False,
                 n_channels=1024,
                 use_quadratic=False, n_bins=8, n_spline_steps=2):
        super(BGAP, self).__init__()
        # assert(n_group_size % 2 == 0)
        self.n_flows = n_flows
        self.n_group_size = n_group_size
        self.transforms = torch.nn.ModuleList()
        self.convinv = torch.nn.ModuleList()
        self.n_speaker_dim = n_speaker_dim
        self.scaling_fn = scaling_fn
        self.attribute_processing = AttributeProcessing(take_log_of_input)
        self.n_spline_steps = n_spline_steps
        self.bottleneck_layer = BottleneckLayerLayer(**bottleneck_hparams)
        n_txt_reduced_dim = self.bottleneck_layer.out_dim
        context_dim = n_txt_reduced_dim * n_group_size + n_speaker_dim

        if self.n_group_size > 1:
            self.unfold_params = {'kernel_size': (n_group_size, 1),
                                  'stride': n_group_size,
                                  'padding': 0, 'dilation': 1}
            self.unfold = nn.Unfold(**self.unfold_params)

        for k in range(n_flows):
            self.convinv.append(Invertible1x1Conv(n_in_dim * n_group_size))
            if k >= n_flows-self.n_spline_steps:
                left = -3
                right= 3
                top = 3
                bottom = -3
                self.transforms.append(SplineTransformationLayer(
                    n_in_dim * n_group_size, context_dim, n_layers,
                    with_dilation=with_dilation, kernel_size=kernel_size,
                    scaling_fn=scaling_fn,
                    n_channels=n_channels, top=top,
                    bottom=bottom, left = left, right=right,
                    use_quadratic=use_quadratic, n_bins=n_bins))
            else:
                self.transforms.append(AffineTransformationLayer(
                    n_in_dim * n_group_size, context_dim, n_layers,
                    with_dilation=with_dilation, kernel_size=kernel_size,
                    scaling_fn=scaling_fn,
                    affine_model='simple_conv', n_channels=n_channels))

    def fold(self, data):
        """Inverse of the self.unfold(data.unsqueeze(-1)) operation used for
        the grouping or "squeeze" operation on input

        Args:
            data: B x C x T tensor of temporal data
        """
        output_size = (data.shape[2]*self.n_group_size, 1)
        data = nn.functional.fold(
            data, output_size=output_size, **self.unfold_params).squeeze(-1)
        return data

    def preprocess_context(self, txt_emb, speaker_vecs, std_scale=None):
        if self.n_group_size > 1:
            txt_emb = self.unfold(txt_emb[..., None])
        speaker_vecs = speaker_vecs[..., None].expand(-1, -1, txt_emb.shape[2])
        context = torch.cat((txt_emb, speaker_vecs), 1)
        return context

    def forward(self, txt_enc, spk_emb, x, lens):
        """x<tensor>: duration or pitch or energy average"""
        assert(txt_enc.size(2) >= x.size(1))
        if len(x.shape) == 2:
            # add channel dimension
            x = x[:, None]
        txt_enc = self.bottleneck_layer(txt_enc)

        # lens including padded values
        lens_grouped = (lens // self.n_group_size).long()
        context = self.preprocess_context(txt_enc, spk_emb)
        x = self.unfold(x[..., None])
        log_s_list, log_det_W_list = [], []
        for k in range(self.n_flows):
            x, log_s = self.transforms[k](x, context, seq_lens=lens_grouped)
            x, log_det_W = self.convinv[k](x)
            log_det_W_list.append(log_det_W)
            log_s_list.append(log_s)
        # prepare outputs
        outputs = {'z': x,
                   'log_det_W_list': log_det_W_list,
                   'log_s_list': log_s_list}

        return outputs

    def infer(self, z, txt_enc, spk_emb, seq_lens):
        txt_enc = self.bottleneck_layer(txt_enc)
        context = self.preprocess_context(txt_enc, spk_emb)
        lens_grouped = (seq_lens // self.n_group_size).long()
        z = self.unfold(z[..., None])
        for k in reversed(range(self.n_flows)):
            z = self.convinv[k](z, inverse=True)
            z = self.transforms[k].forward(z, context,
                                           inverse=True, seq_lens=lens_grouped)
        # z mapped to input domain
        x_hat = self.fold(z)
        # pad on the way out
        return x_hat


class AGAP(torch.nn.Module):
    def __init__(self, n_in_dim, n_speaker_dim, n_flows, n_hidden,
                 n_lstm_layers, bottleneck_hparams, scaling_fn='exp',
                 take_log_of_input=False, p_dropout=0.0, setup='',
                 spline_flow_params=None, n_group_size=1):
        super(AGAP, self).__init__()
        self.flows = torch.nn.ModuleList()
        self.n_group_size = n_group_size
        self.n_speaker_dim = n_speaker_dim
        self.attribute_processing = AttributeProcessing(take_log_of_input)
        self.n_in_dim = n_in_dim
        self.bottleneck_layer = BottleneckLayerLayer(**bottleneck_hparams)
        n_txt_reduced_dim = self.bottleneck_layer.out_dim

        if self.n_group_size > 1:
            self.unfold_params = {'kernel_size': (n_group_size, 1),
                                  'stride': n_group_size,
                                  'padding': 0, 'dilation': 1}
            self.unfold = nn.Unfold(**self.unfold_params)

        if spline_flow_params is not None:
            spline_flow_params['n_in_channels'] *= self.n_group_size

        for i in range(n_flows):
            if i % 2 == 0:
                self.flows.append(AR_Step(
                    n_in_dim * n_group_size, n_speaker_dim, n_txt_reduced_dim *
                    n_group_size, n_hidden, n_lstm_layers, scaling_fn,
                    spline_flow_params))
            else:
                self.flows.append(AR_Back_Step(
                    n_in_dim * n_group_size, n_speaker_dim, n_txt_reduced_dim *
                    n_group_size, n_hidden, n_lstm_layers, scaling_fn,
                    spline_flow_params))

    def fold(self, data):
        """Inverse of the self.unfold(data.unsqueeze(-1)) operation used for
        the grouping or "squeeze" operation on input

        Args:
            data: B x C x T tensor of temporal data
        """
        output_size = (data.shape[2]*self.n_group_size, 1)
        data = nn.functional.fold(
            data, output_size=output_size, **self.unfold_params).squeeze(-1)
        return data

    def preprocess_context(self, txt_emb, speaker_vecs):
        if self.n_group_size > 1:
            txt_emb = self.unfold(txt_emb[..., None])
        speaker_vecs = speaker_vecs[..., None].expand(-1, -1, txt_emb.shape[2])
        context = torch.cat((txt_emb, speaker_vecs), 1)
        return context

    def forward(self, txt_emb, spk_emb, x, lens):
        """x<tensor>: duration or pitch or energy average"""

        x = x[:, None] if len(x.shape) == 2 else x  # add channel dimension
        if self.n_group_size > 1:
            x = self.unfold(x[..., None])
        x = x.permute(2, 0, 1)  # permute to time, batch, dims
        x = self.attribute_processing.normalize(x)

        txt_emb = self.bottleneck_layer(txt_emb)
        context = self.preprocess_context(txt_emb, spk_emb)
        context = context.permute(2, 0, 1)  # permute to time, batch, dims

        lens_groupped = (lens / self.n_group_size).long()
        log_s_list = []
        for i, flow in enumerate(self.flows):
            x, log_s = flow(x, context, lens_groupped)
            log_s_list.append(log_s)

        x = x.permute(1, 2, 0)  # x mapped to z
        log_s_list = [log_s_elt.permute(1, 2, 0) for log_s_elt in log_s_list]
        outputs = {'z': x, 'log_s_list': log_s_list, 'log_det_W_list': []}
        return outputs

    def infer(self, z, txt_emb, spk_emb, seq_lens=None):
        if self.n_group_size > 1:
            n_frames = z.shape[2]
            z = self.unfold(z[..., None])
        z = z.permute(2, 0, 1)  # permute to time, batch, dims

        txt_emb = self.bottleneck_layer(txt_emb)
        context = self.preprocess_context(txt_emb, spk_emb)
        context = context.permute(2, 0, 1)  # permute to time, batch, dims

        for i, flow in enumerate(reversed(self.flows)):
            z = flow.infer(z, context)

        x_hat = z.permute(1, 2, 0)
        if self.n_group_size > 1:
            x_hat = self.fold(x_hat)
            if n_frames > x_hat.shape[2]:
                m = nn.ReflectionPad1d((0, n_frames - x_hat.shape[2]))
                x_hat = m(x_hat)

        x_hat = self.attribute_processing.denormalize(x_hat)
        return x_hat