decoder_base.py

import math
import torch
import torch.nn as nn
from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present

URLS = {
    "hubert-discrete": "https://github.com/bshall/acoustic-model/releases/download/v0.1/hubert-discrete-d49e1c77.pt",
    "hubert-soft": "https://github.com/bshall/acoustic-model/releases/download/v0.1/hubert-soft-0321fd7e.pt",
}

class CustomLSTM(nn.Module):
    def __init__(self, input_sz, hidden_sz):
        super().__init__()
        self.input_sz = input_sz
        self.hidden_size = hidden_sz
        self.W = nn.Parameter(torch.Tensor(input_sz, hidden_sz * 4))
        self.U = nn.Parameter(torch.Tensor(hidden_sz, hidden_sz * 4))
        self.bias = nn.Parameter(torch.Tensor(hidden_sz * 4))
        self.init_weights()

    def init_weights(self):
        stdv = 1.0 / math.sqrt(self.hidden_size)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, stdv)

    def forward(self, x,
                init_states=None):
        """Assumes x is of shape (batch, sequence, feature)"""
        #print(type(x))
        #print(x.shape)
        bs, seq_sz, _ = x.size()
        hidden_seq = []
        if init_states is None:
            h_t, c_t = (torch.zeros(bs, self.hidden_size).to(x.device),
                        torch.zeros(bs, self.hidden_size).to(x.device))
        else:
            h_t, c_t = init_states

        HS = self.hidden_size
        for t in range(seq_sz):
            x_t = x[:, t, :]
            # batch the computations into a single matrix multiplication
            gates = x_t @ self.W + h_t @ self.U + self.bias
            i_t, f_t, g_t, o_t = (
                torch.sigmoid(gates[:, :HS]), # input
                torch.sigmoid(gates[:, HS:HS*2]), # forget
                torch.tanh(gates[:, HS*2:HS*3]),
                torch.sigmoid(gates[:, HS*3:]), # output
            )
            c_t = f_t * c_t + i_t * g_t
            h_t = o_t * torch.tanh(c_t)
            hidden_seq.append(h_t.unsqueeze(0))
        hidden_seq = torch.cat(hidden_seq, dim=0)
        # reshape from shape (sequence, batch, feature) to (batch, sequence, feature)
        hidden_seq = hidden_seq.transpose(0, 1).contiguous()
        return hidden_seq, (h_t, c_t)

class AcousticModel(nn.Module):
    def __init__(self, discrete: bool = False, upsample: bool = True, use_custom_lstm=False):
        super().__init__()
        # self.spk_projection = nn.Linear(512+512, 512)
        self.encoder = Encoder(discrete, upsample)
        self.decoder = Decoder(use_custom_lstm=use_custom_lstm)

    def forward(self, x: torch.Tensor, spk_embs, mels: torch.Tensor) -> torch.Tensor:
        x = self.encoder(x)
        exp_spk_embs = spk_embs.unsqueeze(1).expand(-1, x.size(1), -1)
        concat_x = torch.cat([x, exp_spk_embs], dim=-1)
        # x = self.spk_projection(concat_x)
        return self.decoder(concat_x, mels)

    #def forward(self, x: torch.Tensor, mels: torch.Tensor) -> torch.Tensor:
    #    x = self.encoder(x)
    #    return self.decoder(x, mels)

    def forward_test(self, x, spk_embs, mels):
      print('x shape', x.shape)
      print('se shape', spk_embs.shape)
      print('mels shape', mels.shape)
      x = self.encoder(x)
      print('x_enc shape', x.shape)
      return

    @torch.inference_mode()
    def generate(self, x: torch.Tensor, spk_embs) -> torch.Tensor:
        x = self.encoder(x)
        exp_spk_embs = spk_embs.unsqueeze(1).expand(-1, x.size(1), -1)
        concat_x = torch.cat([x, exp_spk_embs], dim=-1)
        # x = self.spk_projection(concat_x)
        return self.decoder.generate(concat_x)


class Encoder(nn.Module):
    def __init__(self, discrete: bool = False, upsample: bool = True):
        super().__init__()
        self.embedding = nn.Embedding(100 + 1, 256) if discrete else None
        self.prenet = PreNet(256, 256, 256)
        self.convs = nn.Sequential(
            nn.Conv1d(256, 512, 5, 1, 2),
            nn.ReLU(),
            nn.InstanceNorm1d(512),
            nn.ConvTranspose1d(512, 512, 4, 2, 1) if upsample else nn.Identity(),
            nn.Conv1d(512, 512, 5, 1, 2),
            nn.ReLU(),
            nn.InstanceNorm1d(512),
            nn.Conv1d(512, 512, 5, 1, 2),
            nn.ReLU(),
            nn.InstanceNorm1d(512),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.embedding is not None:
            x = self.embedding(x)
        x = self.prenet(x)
        x = self.convs(x.transpose(1, 2))
        return x.transpose(1, 2)


class Decoder(nn.Module):
    def __init__(self, use_custom_lstm=False):
        super().__init__()
        self.use_custom_lstm = use_custom_lstm
        self.prenet = PreNet(128, 256, 256)
        self.prenet = PreNet(128, 256, 256)
        if use_custom_lstm:
          self.lstm1 = CustomLSTM(1024 + 256, 768)
          self.lstm2 = CustomLSTM(768, 768)
          self.lstm3 = CustomLSTM(768, 768)
        else:
          self.lstm1 = nn.LSTM(1024 + 256, 768)
          self.lstm2 = nn.LSTM(768, 768)
          self.lstm3 = nn.LSTM(768, 768)
        self.proj = nn.Linear(768, 128, bias=False)

    def forward(self, x: torch.Tensor, mels: torch.Tensor) -> torch.Tensor:
        mels = self.prenet(mels)
        x, _ = self.lstm1(torch.cat((x, mels), dim=-1))
        res = x
        x, _ = self.lstm2(x)
        x = res + x
        res = x
        x, _ = self.lstm3(x)
        x = res + x
        return self.proj(x)

    @torch.inference_mode()
    def generate(self, xs: torch.Tensor) -> torch.Tensor:
        m = torch.zeros(xs.size(0), 128, device=xs.device)
        if not self.use_custom_lstm:
          h1 = torch.zeros(1, xs.size(0), 768, device=xs.device)
          c1 = torch.zeros(1, xs.size(0), 768, device=xs.device)
          h2 = torch.zeros(1, xs.size(0), 768, device=xs.device)
          c2 = torch.zeros(1, xs.size(0), 768, device=xs.device)
          h3 = torch.zeros(1, xs.size(0), 768, device=xs.device)
          c3 = torch.zeros(1, xs.size(0), 768, device=xs.device)
        else:
          h1 = torch.zeros(xs.size(0), 768, device=xs.device)
          c1 = torch.zeros(xs.size(0), 768, device=xs.device)
          h2 = torch.zeros(xs.size(0), 768, device=xs.device)
          c2 = torch.zeros(xs.size(0), 768, device=xs.device)
          h3 = torch.zeros(xs.size(0), 768, device=xs.device)
          c3 = torch.zeros(xs.size(0), 768, device=xs.device)

        mel = []
        for x in torch.unbind(xs, dim=1):
            m = self.prenet(m)
            x = torch.cat((x, m), dim=1).unsqueeze(1)
            x1, (h1, c1) = self.lstm1(x, (h1, c1))
            x2, (h2, c2) = self.lstm2(x1, (h2, c2))
            x = x1 + x2
            x3, (h3, c3) = self.lstm3(x, (h3, c3))
            x = x + x3
            m = self.proj(x).squeeze(1)
            mel.append(m)
        return torch.stack(mel, dim=1)


class PreNet(nn.Module):
    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        output_size: int,
        dropout: float = 0.5,
    ):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size, output_size),
            nn.ReLU(),
            nn.Dropout(dropout),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


def _acoustic(
    name: str,
    discrete: bool,
    upsample: bool,
    pretrained: bool = True,
    progress: bool = True,
) -> AcousticModel:
    acoustic = AcousticModel(discrete, upsample)
    if pretrained:
        checkpoint = torch.hub.load_state_dict_from_url(URLS[name], progress=progress)
        consume_prefix_in_state_dict_if_present(checkpoint["acoustic-model"], "module.")
        acoustic.load_state_dict(checkpoint["acoustic-model"])
        acoustic.eval()
    return acoustic


def hubert_discrete(
    pretrained: bool = True,
    progress: bool = True,
) -> AcousticModel:
    r"""HuBERT-Discrete acoustic model from `"A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion"`.
    Args:
        pretrained (bool): load pretrained weights into the model
        progress (bool): show progress bar when downloading model
    """
    return _acoustic(
        "hubert-discrete",
        discrete=True,
        upsample=True,
        pretrained=pretrained,
        progress=progress,
    )


def hubert_soft(
    pretrained: bool = True,
    progress: bool = True,
) -> AcousticModel:
    r"""HuBERT-Soft acoustic model from `"A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion"`.
    Args:
        pretrained (bool): load pretrained weights into the model
        progress (bool): show progress bar when downloading model
    """
    return _acoustic(
        "hubert-soft",
        discrete=False,
        upsample=True,
        pretrained=pretrained,
        progress=progress,
    )