Modules/hifigan.py

import torch
import torch.nn.functional as F
import torch.nn as nn
from torch.nn import Conv1d, ConvTranspose1d
from torch.nn.utils.parametrizations import weight_norm
import math
import numpy as np





def get_padding(kernel_size, dilation=1):
    return int((kernel_size*dilation - dilation)/2)


def _tile(x,
          length=None):
    x = x.repeat(1, 1, int(length / x.shape[2]) + 1)[:, :, :length]
    return x


class AdaIN1d(nn.Module):

    # used by HiFiGan & ProsodyPredictor

    def __init__(self, style_dim, num_features):
        super().__init__()
        self.norm = nn.InstanceNorm1d(num_features, affine=False)
        self.fc = nn.Linear(style_dim, num_features*2)

    def forward(self, x, s):

        # x = torch.Size([1, 512, 248])     same as output
        # s = torch.Size([1, 7, 1, 128])

        s = self.fc(s.transpose(1, 2)).transpose(1, 2)

        s = _tile(s, length=x.shape[2])

        gamma, beta = torch.chunk(s, chunks=2, dim=1)
        return (1+gamma) * self.norm(x) + beta


class AdaINResBlock1(torch.nn.Module):
    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), style_dim=64):
        super(AdaINResBlock1, self).__init__()
        self.convs1 = nn.ModuleList([
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
                               padding=get_padding(kernel_size, dilation[0]))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
                               padding=get_padding(kernel_size, dilation[1]))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
                               padding=get_padding(kernel_size, dilation[2])))
        ])
        # self.convs1.apply(init_weights)

        self.convs2 = nn.ModuleList([
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1)))
        ])
        # self.convs2.apply(init_weights)

        self.adain1 = nn.ModuleList([
            AdaIN1d(style_dim, channels),
            AdaIN1d(style_dim, channels),
            AdaIN1d(style_dim, channels),
        ])

        self.adain2 = nn.ModuleList([
            AdaIN1d(style_dim, channels),
            AdaIN1d(style_dim, channels),
            AdaIN1d(style_dim, channels),
        ])

        self.alpha1 = nn.ParameterList(
            [nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs1))])
        self.alpha2 = nn.ParameterList(
            [nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs2))])

    def forward(self, x, s):
        for c1, c2, n1, n2, a1, a2 in zip(self.convs1, self.convs2, self.adain1, self.adain2, self.alpha1, self.alpha2):
            xt = n1(x, s)  # THIS IS ADAIN - EXPECTS conv1d dims
            xt = xt + (1 / a1) * (torch.sin(a1 * xt) ** 2)  # Snake1D
            xt = c1(xt)
            xt = n2(xt, s)  # THIS IS ADAIN - EXPECTS conv1d dims
            xt = xt + (1 / a2) * (torch.sin(a2 * xt) ** 2)  # Snake1D
            xt = c2(xt)
            x = xt + x
        return x


class SourceModuleHnNSF(torch.nn.Module):

    def __init__(self):

        super().__init__()
        self.harmonic_num = 8
        self.l_linear = torch.nn.Linear(self.harmonic_num + 1, 1)
        self.upsample_scale = 300
        

    def forward(self, x):
        # --
        x = torch.multiply(x, torch.FloatTensor(
            [[range(1, self.harmonic_num + 2)]]).to(x.device))  # [1, 145200, 9]
        
        # modulo of negative f0_values => -21 % 10 = 9 as -3*10 + 9 = 21 NOTICE THAT f0_values IS SIGNED
        rad_values = x / 25647 #).clamp(0, 1)
        # rad_values = torch.where(torch.logical_or(rad_values < 0, rad_values > 1), 0.5, rad_values)
        rad_values = rad_values % 1  # % of neg values
        rad_values = F.interpolate(rad_values.transpose(1, 2),
                                                     scale_factor=1/self.upsample_scale,
                                                     mode='linear').transpose(1, 2)

        # 1.89 sounds also nice has woofer at punctuation
        phase = torch.cumsum(rad_values, dim=1) * 1.84 * np.pi
        phase = F.interpolate(phase.transpose(1, 2) * self.upsample_scale,
                              scale_factor=self.upsample_scale, mode='linear').transpose(1, 2)
        x = .009 * phase.sin()
        # --
        x = self.l_linear(x).tanh()
        return x


class Generator(torch.nn.Module):
    def __init__(self,
                 style_dim,
                 resblock_kernel_sizes,
                 upsample_rates,
                 upsample_initial_channel,
                 resblock_dilation_sizes,
                 upsample_kernel_sizes):
        super(Generator, self).__init__()
        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_rates)
        self.m_source = SourceModuleHnNSF()
        self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates))
        self.noise_convs = nn.ModuleList()
        self.ups = nn.ModuleList()
        self.noise_res = nn.ModuleList()

        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            c_cur = upsample_initial_channel // (2 ** (i + 1))

            self.ups.append(weight_norm(ConvTranspose1d(upsample_initial_channel//(2**i),
                                                        upsample_initial_channel//(
                                                            2**(i+1)),
                                                        k, u, padding=(u//2 + u % 2), output_padding=u % 2)))

            if i + 1 < len(upsample_rates):
                stride_f0 = np.prod(upsample_rates[i + 1:])
                self.noise_convs.append(Conv1d(
                    1, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=(stride_f0+1) // 2))
                self.noise_res.append(AdaINResBlock1(
                    c_cur, 7, [1, 3, 5], style_dim))
            else:
                self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
                self.noise_res.append(AdaINResBlock1(
                    c_cur, 11, [1, 3, 5], style_dim))

        self.resblocks = nn.ModuleList()

        self.alphas = nn.ParameterList()
        self.alphas.append(nn.Parameter(
            torch.ones(1, upsample_initial_channel, 1)))

        for i in range(len(self.ups)):
            ch = upsample_initial_channel//(2**(i+1))
            self.alphas.append(nn.Parameter(torch.ones(1, ch, 1)))

            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
                self.resblocks.append(AdaINResBlock1(ch, k, d, style_dim))

        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))

    def forward(self, x, s, f0):

        # x.shape=torch.Size([1, 512, 484]) s.shape=torch.Size([1, 1, 1, 128]) f0.shape=torch.Size([1, 484]) GENERAT 249
        f0 = self.f0_upsamp(f0).transpose(1, 2)

        # x.shape=torch.Size([1, 512, 484]) s.shape=torch.Size([1, 1, 1, 128]) f0.shape=torch.Size([1, 145200, 1]) GENERAT 253

        # [1, 145400, 1] f0 enters already upsampled to full wav 24kHz length
        har_source = self.m_source(f0)

        har_source = har_source.transpose(1, 2)

        for i in range(self.num_upsamples):

            x = x + (1 / self.alphas[i]) * (torch.sin(self.alphas[i] * x) ** 2)
            x_source = self.noise_convs[i](har_source)
            x_source = self.noise_res[i](x_source, s)

            x = self.ups[i](x)

            x = x + x_source

            xs = None
            for j in range(self.num_kernels):

                if xs is None:
                    xs = self.resblocks[i*self.num_kernels+j](x, s)
                else:
                    xs += self.resblocks[i*self.num_kernels+j](x, s)
            x = xs / self.num_kernels
        # x = x + (1 / self.alphas[i+1]) * (torch.sin(self.alphas[i+1] * x) ** 2)  # noisy
        x = self.conv_post(x)
        x = torch.tanh(x)

        return x

class AdainResBlk1d(nn.Module):

    # also used in ProsodyPredictor()

    def __init__(self, dim_in, dim_out, style_dim=64, actv=nn.LeakyReLU(0.2),
                 upsample='none', dropout_p=0.0):
        super().__init__()
        self.actv = actv
        self.upsample_type = upsample
        self.upsample = UpSample1d(upsample)
        self.learned_sc = dim_in != dim_out
        self._build_weights(dim_in, dim_out, style_dim)
        if upsample == 'none':
            self.pool = nn.Identity()
        else:
            self.pool = weight_norm(nn.ConvTranspose1d(
                dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1, output_padding=1))

    def _build_weights(self, dim_in, dim_out, style_dim):
        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
        self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1))
        self.norm1 = AdaIN1d(style_dim, dim_in)
        self.norm2 = AdaIN1d(style_dim, dim_out)
        if self.learned_sc:
            self.conv1x1 = weight_norm(
                nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))

    def _shortcut(self, x):
        x = self.upsample(x)
        if self.learned_sc:
            x = self.conv1x1(x)
        return x

    def _residual(self, x, s):
        x = self.norm1(x, s)
        x = self.actv(x)
        x = self.pool(x)
        x = self.conv1(x)
        x = self.norm2(x, s)
        x = self.actv(x)
        x = self.conv2(x)
        return x

    def forward(self, x, s):
        out = self._residual(x, s)
        out = (out + self._shortcut(x)) / math.sqrt(2)
        return out


class UpSample1d(nn.Module):
    def __init__(self, layer_type):
        super().__init__()
        self.layer_type = layer_type

    def forward(self, x):
        if self.layer_type == 'none':
            return x
        else:
            return F.interpolate(x, scale_factor=2, mode='nearest-exact')


class Decoder(nn.Module):
    def __init__(self, dim_in=512, F0_channel=512, style_dim=64, dim_out=80,
                 resblock_kernel_sizes=[3, 7, 11],
                 upsample_rates=[10, 5, 3, 2],
                 upsample_initial_channel=512,
                 resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
                 upsample_kernel_sizes=[20, 10, 6, 4]):
        super().__init__()

        self.decode = nn.ModuleList()

        self.encode = AdainResBlk1d(dim_in + 2, 1024, style_dim)

        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
        self.decode.append(AdainResBlk1d(
            1024 + 2 + 64, 512, style_dim, upsample=True))

        self.F0_conv = weight_norm(
            nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))  # smooth

        self.N_conv = weight_norm(
            nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))

        self.asr_res = nn.Sequential(
            weight_norm(nn.Conv1d(512, 64, kernel_size=1)),
        )

        self.generator = Generator(style_dim, resblock_kernel_sizes, upsample_rates,
                                   upsample_initial_channel, resblock_dilation_sizes, upsample_kernel_sizes)

    def forward(self, asr=None, F0_curve=None, N=None, s=None):


        F0 = self.F0_conv(F0_curve)
        N = self.N_conv(N)


        x = torch.cat([asr, F0, N], axis=1)

        x = self.encode(x, s)

        asr_res = self.asr_res(asr)

        res = True
        for block in self.decode:
            if res:

                x = torch.cat([x, asr_res, F0, N], axis=1)

            x = block(x, s)
            if block.upsample_type != "none":
                res = False

        x = self.generator(x, s, F0_curve)
        return x