Spaces:
Runtime error
Runtime error
File size: 2,564 Bytes
2ccf6b5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import math
import torch.nn as nn
import torch
import torch.nn.functional as F
import pflow.models.components.vits_modules as modules
import pflow.models.components.commons as commons
class Mish(nn.Module):
def forward(self, x):
return x * torch.tanh(F.softplus(x))
class SinusoidalPosEmb(nn.Module):
def __init__(self, dim):
super(SinusoidalPosEmb, self).__init__()
self.dim = dim
def forward(self, x, scale=1000):
if x.ndim < 1:
x = x.unsqueeze(0)
device = x.device
half_dim = self.dim // 2
emb = math.log(10000) / (half_dim - 1)
emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
return emb
class VitsWNDecoder(nn.Module):
def __init__(self,
in_channels,
out_channels,
hidden_channels,
kernel_size,
dilation_rate,
n_layers,
gin_channels=0,
pe_scale=1000
):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.hidden_channels = hidden_channels
self.kernel_size = kernel_size
self.dilation_rate = dilation_rate
self.n_layers = n_layers
self.gin_channels = gin_channels
self.pe_scale = pe_scale
self.time_pos_emb = SinusoidalPosEmb(hidden_channels * 2)
dim = hidden_channels * 2
self.mlp = nn.Sequential(
nn.Linear(dim, dim * 4),
Mish(),
nn.Linear(dim * 4, dim)
)
self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
self.enc = modules.WN(hidden_channels * 2,
kernel_size,
dilation_rate,
n_layers,
gin_channels=gin_channels)
self.proj = nn.Conv1d(hidden_channels * 2, out_channels, 1)
def forward(self, x, x_mask, mu, t, *args, **kwargs):
# x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)),
# 1).to(x.dtype)
t = self.time_pos_emb(t, scale=self.pe_scale)
t = self.mlp(t)
x = self.pre(x) * x_mask
mu = self.pre(mu)
x = torch.cat((x, mu), dim=1)
x = self.enc(x, x_mask, g=t)
stats = self.proj(x) * x_mask
return stats
|