text: # coding: utf-8
  
  import torch
  from torch import nn
  from torch.nn import functional as F
  from torch.autograd import Variable
  import math
  import numpy as np
  
  from .modules import Embedding, Linear, Conv1d, ConvTranspose1d
  from .modules import HighwayConv1d, get_mask_from_lengths
  from .modules import position_encoding_init
  from .deepvoice3 import AttentionLayer
  
  
  class Encoder(nn.Module):
      def __init__(self, n_vocab, embed_dim, channels, kernel_size=3,
                   n_speakers=1, speaker_embed_dim=16, embedding_weight_std=0.01,
                   padding_idx=None, dropout=0.1):
          super(Encoder, self).__init__()
          self.dropout = dropout
  
          # Text input embeddings
          self.embed_tokens = Embedding(
              n_vocab, embed_dim, padding_idx, embedding_weight_std)
  
          E = embed_dim
          D = channels
          self.convnet = nn.Sequential(
              Conv1d(E, 2 * D, kernel_size=1, padding=0, dilation=1, std_mul=1.0),
              nn.ReLU(inplace=True),
              Conv1d(2 * D, 2 * D, kernel_size=1, padding=0, dilation=1, std_mul=2.0),
  
              HighwayConv1d(2 * D, 2 * D, kernel_size=kernel_size, padding=None,
                            dilation=1, std_mul=1.0, dropout=dropout),
              HighwayConv1d(2 * D, 2 * D, kernel_size=kernel_size, padding=None,
                            dilation=3, std_mul=1.0, dropout=dropout),
              HighwayConv1d(2 * D, 2 * D, kernel_size=kernel_size, padding=None,
                            dilation=9, std_mul=1.0, dropout=dropout),
              HighwayConv1d(2 * D, 2 * D, kernel_size=kernel_size, padding=None,
                            dilation=27, std_mul=1.0, dropout=dropout),
  
              HighwayConv1d(2 * D, 2 * D, kernel_size=kernel_size, padding=None,
                            dilation=1, std_mul=1.0, dropout=dropout),
              HighwayConv1d(2 * D, 2 * D, kernel_size=kernel_size, padding=None,
                            dilation=3, std_mul=1.0, dropout=dropout),
              HighwayConv1d(2 * D, 2 * D, kernel_size=kernel_size, padding=None,
                            dilation=9, std_mul=1.0, dropout=dropout),
              HighwayConv1d(2 * D, 2 * D, kernel_size=kernel_size, padding=None,
                            dilation=27, std_mul=1.0, dropout=dropout),
  
              HighwayConv1d(2 * D, 2 * D, kernel_size=kernel_size, padding=None,
                            dilation=1, std_mul=1.0, dropout=dropout),
              HighwayConv1d(2 * D, 2 * D, kernel_size=kernel_size, padding=None,
                            dilation=1, std_mul=1.0, dropout=dropout),
  
              HighwayConv1d(2 * D, 2 * D, kernel_size=1, padding=0,
                            dilation=1, std_mul=1.0, dropout=dropout),
          )
  
      def forward(self, text_sequences, text_positions=None, lengths=None,
                  speaker_embed=None):
          # embed text_sequences
          # (B, T, E)
          x = self.embed_tokens(text_sequences)
  
          x = self.convnet(x.transpose(1, 2)).transpose(1, 2)
  
          # (B, T, D) and (B, T, D)
          keys, values = x.split(x.size(-1) // 2, dim=-1)
  
          return keys, values
  
  
  class Decoder(nn.Module):
      def __init__(self, embed_dim, in_dim=80, r=5, channels=256, kernel_size=3,
                   n_speakers=1, speaker_embed_dim=16,
                   max_positions=512, padding_idx=None,
                   dropout=0.1,
                   use_memory_mask=False,
                   force_monotonic_attention=False,
                   query_position_rate=1.0,
                   key_position_rate=1.29,
                   window_ahead=3,
                   window_backward=1,
                   key_projection=False,
                   value_projection=False,
                   ):
          super(Decoder, self).__init__()
          self.dropout = dropout
          self.in_dim = in_dim
          self.r = r
  
          D = channels
          F = in_dim * r  # should be r = 1 to replicate
          self.audio_encoder_modules = nn.ModuleList([
              Conv1d(F, D, kernel_size=1, padding=0, dilation=1, std_mul=1.0),
              nn.ReLU(inplace=True),
              Conv1d(D, D, kernel_size=1, padding=0, dilation=1, std_mul=2.0),
              nn.ReLU(inplace=True),
              Conv1d(D, D, kernel_size=1, padding=0, dilation=1, std_mul=2.0),
  
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=1, causal=True, std_mul=1.0, dropout=dropout),
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=3, causal=True, std_mul=1.0, dropout=dropout),
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=9, causal=True, std_mul=1.0, dropout=dropout),
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=27, causal=True, std_mul=1.0, dropout=dropout),
  
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=1, causal=True, std_mul=1.0, dropout=dropout),
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=3, causal=True, std_mul=1.0, dropout=dropout),
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=9, causal=True, std_mul=1.0, dropout=dropout),
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=27, causal=True, std_mul=1.0, dropout=dropout),
  
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=3, causal=True, std_mul=1.0, dropout=dropout),
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=3, causal=True, std_mul=1.0, dropout=dropout),
          ])
  
          self.attention = AttentionLayer(D, D, dropout=dropout,
                                          window_ahead=window_ahead,
                                          window_backward=window_backward,
                                          key_projection=key_projection,
                                          value_projection=value_projection)
  
          self.audio_decoder_modules = nn.ModuleList([
              Conv1d(2 * D, D, kernel_size=1, padding=0, dilation=1, std_mul=1.0),
  
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=1, causal=True, std_mul=1.0, dropout=dropout),
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=3, causal=True, std_mul=1.0, dropout=dropout),
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=9, causal=True, std_mul=1.0, dropout=dropout),
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=27, causal=True, std_mul=1.0, dropout=dropout),
  
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=1, causal=True, std_mul=1.0, dropout=dropout),
              HighwayConv1d(D, D, kernel_size=kernel_size, padding=None,
                            dilation=1, causal=True, std_mul=1.0, dropout=dropout),
  
              Conv1d(D, D, kernel_size=1, padding=0, dilation=1, std_mul=1.0),
              nn.ReLU(inplace=True),
              Conv1d(D, D, kernel_size=1, padding=0, dilation=1, std_mul=2.0),
              nn.ReLU(inplace=True),
              Conv1d(D, D, kernel_size=1, padding=0, dilation=1, std_mul=2.0),
              nn.ReLU(inplace=True),
          ])
          self.last_conv = Conv1d(D, F, kernel_size=1, padding=0, dilation=1, std_mul=2.0)
  
          # Done prediction
          self.fc = Linear(F, 1)
  
          # Position encodings for query (decoder states) and keys (encoder states)
          self.embed_query_positions = Embedding(
              max_positions, D, padding_idx)
          self.embed_query_positions.weight.data = position_encoding_init(
              max_positions, D, position_rate=query_position_rate, sinusoidal=True)
          self.embed_keys_positions = Embedding(
              max_positions, D, padding_idx)
          self.embed_keys_positions.weight.data = position_encoding_init(
              max_positions, D, position_rate=key_position_rate, sinusoidal=True)
  
          # options
          self.max_decoder_steps = 200
          self.min_decoder_steps = 10
          self.use_memory_mask = use_memory_mask
          self.force_monotonic_attention = force_monotonic_attention
  
      def forward(self, encoder_out, inputs=None,
                  text_positions=None, frame_positions=None,
                  speaker_embed=None, lengths=None):
  
          if inputs is None:
              assert text_positions is not None
              self.start_fresh_sequence()
              outputs = self.incremental_forward(encoder_out, text_positions)
              return outputs
  
          # Grouping multiple frames if necessary
          if inputs.size(-1) == self.in_dim:
              inputs = inputs.view(inputs.size(0), inputs.size(1) // self.r, -1)
          assert inputs.size(-1) == self.in_dim * self.r
  
          keys, values = encoder_out
  
          if self.use_memory_mask and lengths is not None:
              mask = get_mask_from_lengths(keys, lengths)
          else:
              mask = None
  
          # position encodings
          if text_positions is not None:
              text_pos_embed = self.embed_keys_positions(text_positions)
              keys = keys + text_pos_embed
          if frame_positions is not None:
              frame_pos_embed = self.embed_query_positions(frame_positions)
  
          # transpose only once to speed up attention layers
          keys = keys.transpose(1, 2).contiguous()
  
          # (B, T, C)
          x = inputs
  
          # (B, C, T)
          x = x.transpose(1, 2)
  
          # Apply audio encoder
          for f in self.audio_encoder_modules:
              x = f(x)
          Q = x
  
          # Attention modules assume query as (B, T, C)
          x = x.transpose(1, 2)
          x = x if frame_positions is None else x + frame_pos_embed
          R, alignments = self.attention(x, (keys, values), mask=mask)
          R = R.transpose(1, 2)
  
          # (B, C*2, T)
          Rd = torch.cat((R, Q), dim=1)
          x = Rd
  
          # Apply audio decoder
          for f in self.audio_decoder_modules:
              x = f(x)
          decoder_states = x.transpose(1, 2).contiguous()
          x = self.last_conv(x)
  
          # (B, T, C)
          x = x.transpose(1, 2)
  
          # Mel
          outputs = F.sigmoid(x)
  
          # Done prediction
          done = F.sigmoid(self.fc(x))
  
          # Adding extra dim for convenient
          alignments = alignments.unsqueeze(0)
  
          return outputs, alignments, done, decoder_states
  
      def incremental_forward(self, encoder_out, text_positions,
                              initial_input=None, test_inputs=None):
          keys, values = encoder_out
          B = keys.size(0)
  
          # position encodings
          if text_positions is not None:
              text_pos_embed = self.embed_keys_positions(text_positions)
              keys = keys + text_pos_embed
  
          # transpose only once to speed up attention layers
          keys = keys.transpose(1, 2).contiguous()
  
          decoder_states = []
          outputs = []
          alignments = []
          dones = []
          # intially set to zeros
          last_attended = 0 if self.force_monotonic_attention else None
  
          t = 0
          if initial_input is None:
              initial_input = Variable(
                  keys.data.new(B, 1, self.in_dim * self.r).zero_())
          current_input = initial_input
          while True:
              # frame pos start with 1.
              frame_pos = Variable(keys.data.new(B, 1).fill_(t + 1)).long()
              frame_pos_embed = self.embed_query_positions(frame_pos)
  
              if test_inputs is not None:
                  if t >= test_inputs.size(1):
                      break
                  current_input = test_inputs[:, t, :].unsqueeze(1)
              else:
                  if t > 0:
                      current_input = outputs[-1]
  
              # (B, 1, C)
              x = current_input
  
              for f in self.audio_encoder_modules:
                  try:
                      x = f.incremental_forward(x)
                  except AttributeError as e:
                      x = f(x)
              Q = x
  
              R, alignment = self.attention(
                  x + frame_pos_embed, (keys, values), last_attended=last_attended)
              if self.force_monotonic_attention:
                  last_attended = alignment.max(-1)[1].view(-1).data[0]
  
              Rd = torch.cat((R, Q), dim=-1)
              x = Rd
              for f in self.audio_decoder_modules:
                  try:
                      x = f.incremental_forward(x)
                  except AttributeError as e:
                      x = f(x)
              decoder_state = x
              x = self.last_conv.incremental_forward(x)
  
              # Ooutput & done flag predictions
              output = F.sigmoid(x)
              done = F.sigmoid(self.fc(x))
  
              decoder_states += [decoder_state]
              outputs += [output]
              alignments += [alignment]
              dones += [done]
  
              t += 1
              if test_inputs is None:
                  if (done > 0.5).all() and t > self.min_decoder_steps:
                      break
                  elif t > self.max_decoder_steps:
                      break
  
          # Remove 1-element time axis
          alignments = list(map(lambda x: x.squeeze(1), alignments))
          decoder_states = list(map(lambda x: x.squeeze(1), decoder_states))
          outputs = list(map(lambda x: x.squeeze(1), outputs))
  
          # Combine outputs for all time steps
          alignments = torch.stack(alignments).transpose(0, 1)
          decoder_states = torch.stack(decoder_states).transpose(0, 1).contiguous()
          outputs = torch.stack(outputs).transpose(0, 1).contiguous()
  
          return outputs, alignments, dones, decoder_states
  
      def start_fresh_sequence(self):
          _clear_modules(self.audio_encoder_modules)
          _clear_modules(self.audio_decoder_modules)
          self.last_conv.clear_buffer()
  
  
  def _clear_modules(modules):
      for m in modules:
          try:
              m.clear_buffer()
          except AttributeError as e:
              pass
  
  
  class Converter(nn.Module):
      def __init__(self, in_dim, out_dim, channels=512,  kernel_size=3, dropout=0.1):
          super(Converter, self).__init__()
          self.dropout = dropout
          self.in_dim = in_dim
          self.out_dim = out_dim
  
          F = in_dim
          Fd = out_dim
          C = channels
          self.convnet = nn.Sequential(
              Conv1d(F, C, kernel_size=1, padding=0, dilation=1, std_mul=1.0),
  
              HighwayConv1d(C, C, kernel_size=kernel_size, padding=None,
                            dilation=1, std_mul=1.0, dropout=dropout),
              HighwayConv1d(C, C, kernel_size=kernel_size, padding=None,
                            dilation=3, std_mul=1.0, dropout=dropout),
  
              ConvTranspose1d(C, C, kernel_size=2, padding=0, stride=2, std_mul=1.0),
              HighwayConv1d(C, C, kernel_size=kernel_size, padding=None,
                            dilation=1, std_mul=1.0, dropout=dropout),
              HighwayConv1d(C, C, kernel_size=kernel_size, padding=None,
                            dilation=3, std_mul=1.0, dropout=dropout),
              ConvTranspose1d(C, C, kernel_size=2, padding=0, stride=2, std_mul=1.0),
              HighwayConv1d(C, C, kernel_size=kernel_size, padding=None,
                            dilation=1, std_mul=1.0, dropout=dropout),
              HighwayConv1d(C, C, kernel_size=kernel_size, padding=None,
                            dilation=3, std_mul=1.0, dropout=dropout),
  
              Conv1d(C, 2 * C, kernel_size=1, padding=0, dilation=1, std_mul=1.0),
  
              HighwayConv1d(2 * C, 2 * C, kernel_size=kernel_size, padding=None,
                            dilation=1, std_mul=1.0, dropout=dropout),
              HighwayConv1d(2 * C, 2 * C, kernel_size=kernel_size, padding=None,
                            dilation=1, std_mul=1.0, dropout=dropout),
  
              Conv1d(2 * C, Fd, kernel_size=1, padding=0, dilation=1, std_mul=1.0),
  
              Conv1d(Fd, Fd, kernel_size=1, padding=0, dilation=1, std_mul=1.0),
              nn.ReLU(inplace=True),
              Conv1d(Fd, Fd, kernel_size=1, padding=0, dilation=1, std_mul=2.0),
              nn.ReLU(inplace=True),
  
              Conv1d(Fd, Fd, kernel_size=1, padding=0, dilation=1, std_mul=2.0),
              nn.Sigmoid(),
          )
  
      def forward(self, x, speaker_embed=None):
          return self.convnet(x.transpose(1, 2)).transpose(1, 2)
  