fairseq/models/text_to_speech/tacotron2.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 torch
from torch import nn
from torch.nn import functional as F

from fairseq.models import (
    FairseqEncoder,
    FairseqEncoderDecoderModel,
    FairseqIncrementalDecoder,
    register_model,
    register_model_architecture,
)
from fairseq.modules import LSTMCellWithZoneOut, LocationAttention


logger = logging.getLogger(__name__)


def encoder_init(m):
    if isinstance(m, nn.Conv1d):
        nn.init.xavier_uniform_(m.weight, torch.nn.init.calculate_gain("relu"))


class Tacotron2Encoder(FairseqEncoder):
    def __init__(self, args, src_dict, embed_speaker):
        super().__init__(src_dict)
        self.padding_idx = src_dict.pad()
        self.embed_speaker = embed_speaker
        self.spk_emb_proj = None
        if embed_speaker is not None:
            self.spk_emb_proj = nn.Linear(
                args.encoder_embed_dim + args.speaker_embed_dim, args.encoder_embed_dim
            )

        self.embed_tokens = nn.Embedding(
            len(src_dict), args.encoder_embed_dim, padding_idx=self.padding_idx
        )

        assert args.encoder_conv_kernel_size % 2 == 1
        self.convolutions = nn.ModuleList(
            nn.Sequential(
                nn.Conv1d(
                    args.encoder_embed_dim,
                    args.encoder_embed_dim,
                    kernel_size=args.encoder_conv_kernel_size,
                    padding=((args.encoder_conv_kernel_size - 1) // 2),
                ),
                nn.BatchNorm1d(args.encoder_embed_dim),
                nn.ReLU(),
                nn.Dropout(args.encoder_dropout),
            )
            for _ in range(args.encoder_conv_layers)
        )

        self.lstm = nn.LSTM(
            args.encoder_embed_dim,
            args.encoder_embed_dim // 2,
            num_layers=args.encoder_lstm_layers,
            batch_first=True,
            bidirectional=True,
        )

        self.apply(encoder_init)

    def forward(self, src_tokens, src_lengths=None, speaker=None, **kwargs):
        x = self.embed_tokens(src_tokens)
        x = x.transpose(1, 2).contiguous()  # B x T x C -> B x C x T
        for conv in self.convolutions:
            x = conv(x)
        x = x.transpose(1, 2).contiguous()  # B x C x T -> B x T x C

        src_lengths = src_lengths.cpu().long()
        x = nn.utils.rnn.pack_padded_sequence(x, src_lengths, batch_first=True)
        x = self.lstm(x)[0]
        x = nn.utils.rnn.pad_packed_sequence(x, batch_first=True)[0]

        encoder_padding_mask = src_tokens.eq(self.padding_idx)

        if self.embed_speaker is not None:
            seq_len, bsz, _ = x.size()
            emb = self.embed_speaker(speaker).expand(seq_len, bsz, -1)
            x = self.spk_emb_proj(torch.cat([x, emb], dim=2))

        return {
            "encoder_out": [x],  # B x T x C
            "encoder_padding_mask": encoder_padding_mask,  # B x T
        }


class Prenet(nn.Module):
    def __init__(self, in_dim, n_layers, n_units, dropout):
        super().__init__()
        self.layers = nn.ModuleList(
            nn.Sequential(nn.Linear(in_dim if i == 0 else n_units, n_units), nn.ReLU())
            for i in range(n_layers)
        )
        self.dropout = dropout

    def forward(self, x):
        for layer in self.layers:
            x = F.dropout(layer(x), p=self.dropout)  # always applies dropout
        return x


class Postnet(nn.Module):
    def __init__(self, in_dim, n_channels, kernel_size, n_layers, dropout):
        super(Postnet, self).__init__()
        self.convolutions = nn.ModuleList()
        assert kernel_size % 2 == 1
        for i in range(n_layers):
            cur_layers = (
                [
                    nn.Conv1d(
                        in_dim if i == 0 else n_channels,
                        n_channels if i < n_layers - 1 else in_dim,
                        kernel_size=kernel_size,
                        padding=((kernel_size - 1) // 2),
                    ),
                    nn.BatchNorm1d(n_channels if i < n_layers - 1 else in_dim),
                ]
                + ([nn.Tanh()] if i < n_layers - 1 else [])
                + [nn.Dropout(dropout)]
            )
            nn.init.xavier_uniform_(
                cur_layers[0].weight,
                torch.nn.init.calculate_gain("tanh" if i < n_layers - 1 else "linear"),
            )
            self.convolutions.append(nn.Sequential(*cur_layers))

    def forward(self, x):
        x = x.transpose(1, 2)  # B x T x C -> B x C x T
        for conv in self.convolutions:
            x = conv(x)
        return x.transpose(1, 2)


def decoder_init(m):
    if isinstance(m, torch.nn.Conv1d):
        nn.init.xavier_uniform_(m.weight, torch.nn.init.calculate_gain("tanh"))


class Tacotron2Decoder(FairseqIncrementalDecoder):
    def __init__(self, args, src_dict):
        super().__init__(None)
        self.args = args
        self.n_frames_per_step = args.n_frames_per_step
        self.out_dim = args.output_frame_dim * args.n_frames_per_step

        self.prenet = Prenet(
            self.out_dim, args.prenet_layers, args.prenet_dim, args.prenet_dropout
        )

        # take prev_context, prev_frame, (speaker embedding) as input
        self.attention_lstm = LSTMCellWithZoneOut(
            args.zoneout,
            args.prenet_dim + args.encoder_embed_dim,
            args.decoder_lstm_dim,
        )

        # take attention_lstm output, attention_state, encoder_out as input
        self.attention = LocationAttention(
            args.attention_dim,
            args.encoder_embed_dim,
            args.decoder_lstm_dim,
            (1 + int(args.attention_use_cumprob)),
            args.attention_conv_dim,
            args.attention_conv_kernel_size,
        )

        # take attention_lstm output, context, (gated_latent) as input
        self.lstm = nn.ModuleList(
            LSTMCellWithZoneOut(
                args.zoneout,
                args.encoder_embed_dim + args.decoder_lstm_dim,
                args.decoder_lstm_dim,
            )
            for i in range(args.decoder_lstm_layers)
        )

        proj_in_dim = args.encoder_embed_dim + args.decoder_lstm_dim
        self.feat_proj = nn.Linear(proj_in_dim, self.out_dim)
        self.eos_proj = nn.Linear(proj_in_dim, 1)

        self.postnet = Postnet(
            self.out_dim,
            args.postnet_conv_dim,
            args.postnet_conv_kernel_size,
            args.postnet_layers,
            args.postnet_dropout,
        )

        self.ctc_proj = None
        if getattr(args, "ctc_weight", 0.0) > 0.0:
            self.ctc_proj = nn.Linear(self.out_dim, len(src_dict))

        self.apply(decoder_init)

    def _get_states(self, incremental_state, enc_out):
        bsz, in_len, _ = enc_out.size()
        alstm_h = self.get_incremental_state(incremental_state, "alstm_h")
        if alstm_h is None:
            alstm_h = enc_out.new_zeros(bsz, self.args.decoder_lstm_dim)
        alstm_c = self.get_incremental_state(incremental_state, "alstm_c")
        if alstm_c is None:
            alstm_c = enc_out.new_zeros(bsz, self.args.decoder_lstm_dim)

        lstm_h = self.get_incremental_state(incremental_state, "lstm_h")
        if lstm_h is None:
            lstm_h = [
                enc_out.new_zeros(bsz, self.args.decoder_lstm_dim)
                for _ in range(self.args.decoder_lstm_layers)
            ]
        lstm_c = self.get_incremental_state(incremental_state, "lstm_c")
        if lstm_c is None:
            lstm_c = [
                enc_out.new_zeros(bsz, self.args.decoder_lstm_dim)
                for _ in range(self.args.decoder_lstm_layers)
            ]

        attn_w = self.get_incremental_state(incremental_state, "attn_w")
        if attn_w is None:
            attn_w = enc_out.new_zeros(bsz, in_len)
        attn_w_cum = self.get_incremental_state(incremental_state, "attn_w_cum")
        if attn_w_cum is None:
            attn_w_cum = enc_out.new_zeros(bsz, in_len)
        return alstm_h, alstm_c, lstm_h, lstm_c, attn_w, attn_w_cum

    def _get_init_attn_c(self, enc_out, enc_mask):
        bsz = enc_out.size(0)
        if self.args.init_attn_c == "zero":
            return enc_out.new_zeros(bsz, self.args.encoder_embed_dim)
        elif self.args.init_attn_c == "avg":
            enc_w = (~enc_mask).type(enc_out.type())
            enc_w = enc_w / enc_w.sum(dim=1, keepdim=True)
            return torch.sum(enc_out * enc_w.unsqueeze(2), dim=1)
        else:
            raise ValueError(f"{self.args.init_attn_c} not supported")

    def forward(
        self,
        prev_output_tokens,
        encoder_out=None,
        incremental_state=None,
        target_lengths=None,
        **kwargs,
    ):
        enc_mask = encoder_out["encoder_padding_mask"]
        enc_out = encoder_out["encoder_out"][0]
        in_len = enc_out.size(1)

        if incremental_state is not None:
            prev_output_tokens = prev_output_tokens[:, -1:, :]
        bsz, out_len, _ = prev_output_tokens.size()

        prenet_out = self.prenet(prev_output_tokens)
        (alstm_h, alstm_c, lstm_h, lstm_c, attn_w, attn_w_cum) = self._get_states(
            incremental_state, enc_out
        )
        attn_ctx = self._get_init_attn_c(enc_out, enc_mask)

        attn_out = enc_out.new_zeros(bsz, in_len, out_len)
        feat_out = enc_out.new_zeros(bsz, out_len, self.out_dim)
        eos_out = enc_out.new_zeros(bsz, out_len)
        for t in range(out_len):
            alstm_in = torch.cat((attn_ctx, prenet_out[:, t, :]), dim=1)
            alstm_h, alstm_c = self.attention_lstm(alstm_in, (alstm_h, alstm_c))

            attn_state = attn_w.unsqueeze(1)
            if self.args.attention_use_cumprob:
                attn_state = torch.stack((attn_w, attn_w_cum), dim=1)
            attn_ctx, attn_w = self.attention(enc_out, enc_mask, alstm_h, attn_state)
            attn_w_cum = attn_w_cum + attn_w
            attn_out[:, :, t] = attn_w

            for i, cur_lstm in enumerate(self.lstm):
                if i == 0:
                    lstm_in = torch.cat((attn_ctx, alstm_h), dim=1)
                else:
                    lstm_in = torch.cat((attn_ctx, lstm_h[i - 1]), dim=1)
                lstm_h[i], lstm_c[i] = cur_lstm(lstm_in, (lstm_h[i], lstm_c[i]))

            proj_in = torch.cat((attn_ctx, lstm_h[-1]), dim=1)
            feat_out[:, t, :] = self.feat_proj(proj_in)
            eos_out[:, t] = self.eos_proj(proj_in).squeeze(1)
        self.attention.clear_cache()

        self.set_incremental_state(incremental_state, "alstm_h", alstm_h)
        self.set_incremental_state(incremental_state, "alstm_c", alstm_c)
        self.set_incremental_state(incremental_state, "lstm_h", lstm_h)
        self.set_incremental_state(incremental_state, "lstm_c", lstm_c)
        self.set_incremental_state(incremental_state, "attn_w", attn_w)
        self.set_incremental_state(incremental_state, "attn_w_cum", attn_w_cum)

        post_feat_out = feat_out + self.postnet(feat_out)
        eos_out = eos_out.view(bsz, out_len, 1)
        return post_feat_out, eos_out, {"attn": attn_out, "feature_out": feat_out}


@register_model("tacotron_2")
class Tacotron2Model(FairseqEncoderDecoderModel):
    """
    Implementation for https://arxiv.org/pdf/1712.05884.pdf
    """

    @staticmethod
    def add_args(parser):
        # encoder
        parser.add_argument("--encoder-dropout", type=float)
        parser.add_argument("--encoder-embed-dim", type=int)
        parser.add_argument("--encoder-conv-layers", type=int)
        parser.add_argument("--encoder-conv-kernel-size", type=int)
        parser.add_argument("--encoder-lstm-layers", type=int)
        # decoder
        parser.add_argument("--attention-dim", type=int)
        parser.add_argument("--attention-conv-dim", type=int)
        parser.add_argument("--attention-conv-kernel-size", type=int)
        parser.add_argument("--prenet-dropout", type=float)
        parser.add_argument("--prenet-layers", type=int)
        parser.add_argument("--prenet-dim", type=int)
        parser.add_argument("--postnet-dropout", type=float)
        parser.add_argument("--postnet-layers", type=int)
        parser.add_argument("--postnet-conv-dim", type=int)
        parser.add_argument("--postnet-conv-kernel-size", type=int)
        parser.add_argument("--init-attn-c", type=str)
        parser.add_argument("--attention-use-cumprob", action="store_true")
        parser.add_argument("--zoneout", type=float)
        parser.add_argument("--decoder-lstm-layers", type=int)
        parser.add_argument("--decoder-lstm-dim", type=int)
        parser.add_argument("--output-frame-dim", type=int)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._num_updates = 0

    @classmethod
    def build_model(cls, args, task):
        embed_speaker = task.get_speaker_embeddings(args)
        encoder = Tacotron2Encoder(args, task.src_dict, embed_speaker)
        decoder = Tacotron2Decoder(args, task.src_dict)
        return cls(encoder, decoder)

    def forward_encoder(self, src_tokens, src_lengths, **kwargs):
        return self.encoder(src_tokens, src_lengths=src_lengths, **kwargs)

    def set_num_updates(self, num_updates):
        super().set_num_updates(num_updates)
        self._num_updates = num_updates


@register_model_architecture("tacotron_2", "tacotron_2")
def base_architecture(args):
    # encoder
    args.encoder_dropout = getattr(args, "encoder_dropout", 0.5)
    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 512)
    args.encoder_conv_layers = getattr(args, "encoder_conv_layers", 3)
    args.encoder_conv_kernel_size = getattr(args, "encoder_conv_kernel_size", 5)
    args.encoder_lstm_layers = getattr(args, "encoder_lstm_layers", 1)
    # decoder
    args.attention_dim = getattr(args, "attention_dim", 128)
    args.attention_conv_dim = getattr(args, "attention_conv_dim", 32)
    args.attention_conv_kernel_size = getattr(args, "attention_conv_kernel_size", 15)
    args.prenet_dropout = getattr(args, "prenet_dropout", 0.5)
    args.prenet_layers = getattr(args, "prenet_layers", 2)
    args.prenet_dim = getattr(args, "prenet_dim", 256)
    args.postnet_dropout = getattr(args, "postnet_dropout", 0.5)
    args.postnet_layers = getattr(args, "postnet_layers", 5)
    args.postnet_conv_dim = getattr(args, "postnet_conv_dim", 512)
    args.postnet_conv_kernel_size = getattr(args, "postnet_conv_kernel_size", 5)
    args.init_attn_c = getattr(args, "init_attn_c", "zero")
    args.attention_use_cumprob = getattr(args, "attention_use_cumprob", True)
    args.zoneout = getattr(args, "zoneout", 0.1)
    args.decoder_lstm_layers = getattr(args, "decoder_lstm_layers", 2)
    args.decoder_lstm_dim = getattr(args, "decoder_lstm_dim", 1024)
    args.output_frame_dim = getattr(args, "output_frame_dim", 80)