mirror of
https://github.com/AIGC-Audio/AudioGPT.git
synced 2026-09-02 04:00:43 +02:00
768 lines
28 KiB
Python
768 lines
28 KiB
Python
import scipy
|
|
from torch.nn import functional as F
|
|
import torch
|
|
from torch import nn
|
|
import numpy as np
|
|
from modules.commons.common_layers import Permute
|
|
from modules.fastspeech.tts_modules import FFTBlocks
|
|
from modules.GenerSpeech.model.wavenet import fused_add_tanh_sigmoid_multiply, WN
|
|
|
|
|
|
class LayerNorm(nn.Module):
|
|
def __init__(self, channels, eps=1e-4):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.eps = eps
|
|
|
|
self.gamma = nn.Parameter(torch.ones(channels))
|
|
self.beta = nn.Parameter(torch.zeros(channels))
|
|
|
|
def forward(self, x):
|
|
n_dims = len(x.shape)
|
|
mean = torch.mean(x, 1, keepdim=True)
|
|
variance = torch.mean((x - mean) ** 2, 1, keepdim=True)
|
|
|
|
x = (x - mean) * torch.rsqrt(variance + self.eps)
|
|
|
|
shape = [1, -1] + [1] * (n_dims - 2)
|
|
x = x * self.gamma.view(*shape) + self.beta.view(*shape)
|
|
return x
|
|
|
|
|
|
class ConvReluNorm(nn.Module):
|
|
def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):
|
|
super().__init__()
|
|
self.in_channels = in_channels
|
|
self.hidden_channels = hidden_channels
|
|
self.out_channels = out_channels
|
|
self.kernel_size = kernel_size
|
|
self.n_layers = n_layers
|
|
self.p_dropout = p_dropout
|
|
assert n_layers > 1, "Number of layers should be larger than 0."
|
|
|
|
self.conv_layers = nn.ModuleList()
|
|
self.norm_layers = nn.ModuleList()
|
|
self.conv_layers.append(nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size // 2))
|
|
self.norm_layers.append(LayerNorm(hidden_channels))
|
|
self.relu_drop = nn.Sequential(
|
|
nn.ReLU(),
|
|
nn.Dropout(p_dropout))
|
|
for _ in range(n_layers - 1):
|
|
self.conv_layers.append(nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size // 2))
|
|
self.norm_layers.append(LayerNorm(hidden_channels))
|
|
self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
|
|
self.proj.weight.data.zero_()
|
|
self.proj.bias.data.zero_()
|
|
|
|
def forward(self, x, x_mask):
|
|
x_org = x
|
|
for i in range(self.n_layers):
|
|
x = self.conv_layers[i](x * x_mask)
|
|
x = self.norm_layers[i](x)
|
|
x = self.relu_drop(x)
|
|
x = x_org + self.proj(x)
|
|
return x * x_mask
|
|
|
|
|
|
|
|
class ActNorm(nn.Module): # glow中的线性变换层
|
|
def __init__(self, channels, ddi=False, **kwargs):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.initialized = not ddi
|
|
|
|
self.logs = nn.Parameter(torch.zeros(1, channels, 1))
|
|
self.bias = nn.Parameter(torch.zeros(1, channels, 1))
|
|
|
|
def forward(self, x, x_mask=None, reverse=False, **kwargs):
|
|
if x_mask is None:
|
|
x_mask = torch.ones(x.size(0), 1, x.size(2)).to(device=x.device, dtype=x.dtype)
|
|
x_len = torch.sum(x_mask, [1, 2])
|
|
if not self.initialized:
|
|
self.initialize(x, x_mask)
|
|
self.initialized = True
|
|
|
|
if reverse:
|
|
z = (x - self.bias) * torch.exp(-self.logs) * x_mask
|
|
logdet = torch.sum(-self.logs) * x_len
|
|
else:
|
|
z = (self.bias + torch.exp(self.logs) * x) * x_mask
|
|
logdet = torch.sum(self.logs) * x_len # [b]
|
|
return z, logdet
|
|
|
|
def store_inverse(self):
|
|
pass
|
|
|
|
def set_ddi(self, ddi):
|
|
self.initialized = not ddi
|
|
|
|
def initialize(self, x, x_mask):
|
|
with torch.no_grad():
|
|
denom = torch.sum(x_mask, [0, 2])
|
|
m = torch.sum(x * x_mask, [0, 2]) / denom
|
|
m_sq = torch.sum(x * x * x_mask, [0, 2]) / denom
|
|
v = m_sq - (m ** 2)
|
|
logs = 0.5 * torch.log(torch.clamp_min(v, 1e-6))
|
|
|
|
bias_init = (-m * torch.exp(-logs)).view(*self.bias.shape).to(dtype=self.bias.dtype)
|
|
logs_init = (-logs).view(*self.logs.shape).to(dtype=self.logs.dtype)
|
|
|
|
self.bias.data.copy_(bias_init)
|
|
self.logs.data.copy_(logs_init)
|
|
|
|
|
|
class InvConvNear(nn.Module): # 可逆卷积
|
|
def __init__(self, channels, n_split=4, no_jacobian=False, lu=True, n_sqz=2, **kwargs):
|
|
super().__init__()
|
|
assert (n_split % 2 == 0)
|
|
self.channels = channels
|
|
self.n_split = n_split
|
|
self.n_sqz = n_sqz
|
|
self.no_jacobian = no_jacobian
|
|
|
|
w_init = torch.qr(torch.FloatTensor(self.n_split, self.n_split).normal_())[0]
|
|
if torch.det(w_init) < 0:
|
|
w_init[:, 0] = -1 * w_init[:, 0]
|
|
self.lu = lu
|
|
if lu:
|
|
# LU decomposition can slightly speed up the inverse
|
|
np_p, np_l, np_u = scipy.linalg.lu(w_init)
|
|
np_s = np.diag(np_u)
|
|
np_sign_s = np.sign(np_s)
|
|
np_log_s = np.log(np.abs(np_s))
|
|
np_u = np.triu(np_u, k=1)
|
|
l_mask = np.tril(np.ones(w_init.shape, dtype=float), -1)
|
|
eye = np.eye(*w_init.shape, dtype=float)
|
|
|
|
self.register_buffer('p', torch.Tensor(np_p.astype(float)))
|
|
self.register_buffer('sign_s', torch.Tensor(np_sign_s.astype(float)))
|
|
self.l = nn.Parameter(torch.Tensor(np_l.astype(float)), requires_grad=True)
|
|
self.log_s = nn.Parameter(torch.Tensor(np_log_s.astype(float)), requires_grad=True)
|
|
self.u = nn.Parameter(torch.Tensor(np_u.astype(float)), requires_grad=True)
|
|
self.register_buffer('l_mask', torch.Tensor(l_mask))
|
|
self.register_buffer('eye', torch.Tensor(eye))
|
|
else:
|
|
self.weight = nn.Parameter(w_init)
|
|
|
|
def forward(self, x, x_mask=None, reverse=False, **kwargs):
|
|
b, c, t = x.size()
|
|
assert (c % self.n_split == 0)
|
|
if x_mask is None:
|
|
x_mask = 1
|
|
x_len = torch.ones((b,), dtype=x.dtype, device=x.device) * t
|
|
else:
|
|
x_len = torch.sum(x_mask, [1, 2])
|
|
|
|
x = x.view(b, self.n_sqz, c // self.n_split, self.n_split // self.n_sqz, t)
|
|
x = x.permute(0, 1, 3, 2, 4).contiguous().view(b, self.n_split, c // self.n_split, t)
|
|
|
|
if self.lu:
|
|
self.weight, log_s = self._get_weight()
|
|
logdet = log_s.sum()
|
|
logdet = logdet * (c / self.n_split) * x_len
|
|
else:
|
|
logdet = torch.logdet(self.weight) * (c / self.n_split) * x_len # [b]
|
|
|
|
if reverse:
|
|
if hasattr(self, "weight_inv"):
|
|
weight = self.weight_inv
|
|
else:
|
|
weight = torch.inverse(self.weight.float()).to(dtype=self.weight.dtype)
|
|
logdet = -logdet
|
|
else:
|
|
weight = self.weight
|
|
if self.no_jacobian:
|
|
logdet = 0
|
|
|
|
weight = weight.view(self.n_split, self.n_split, 1, 1)
|
|
z = F.conv2d(x, weight)
|
|
|
|
z = z.view(b, self.n_sqz, self.n_split // self.n_sqz, c // self.n_split, t)
|
|
z = z.permute(0, 1, 3, 2, 4).contiguous().view(b, c, t) * x_mask
|
|
return z, logdet
|
|
|
|
def _get_weight(self):
|
|
l, log_s, u = self.l, self.log_s, self.u
|
|
l = l * self.l_mask + self.eye
|
|
u = u * self.l_mask.transpose(0, 1).contiguous() + torch.diag(self.sign_s * torch.exp(log_s))
|
|
weight = torch.matmul(self.p, torch.matmul(l, u))
|
|
return weight, log_s
|
|
|
|
def store_inverse(self):
|
|
weight, _ = self._get_weight()
|
|
self.weight_inv = torch.inverse(weight.float()).to(next(self.parameters()).device)
|
|
|
|
|
|
class InvConv(nn.Module):
|
|
def __init__(self, channels, no_jacobian=False, lu=True, **kwargs):
|
|
super().__init__()
|
|
w_shape = [channels, channels]
|
|
w_init = np.linalg.qr(np.random.randn(*w_shape))[0].astype(float)
|
|
LU_decomposed = lu
|
|
if not LU_decomposed:
|
|
# Sample a random orthogonal matrix:
|
|
self.register_parameter("weight", nn.Parameter(torch.Tensor(w_init)))
|
|
else:
|
|
np_p, np_l, np_u = scipy.linalg.lu(w_init)
|
|
np_s = np.diag(np_u)
|
|
np_sign_s = np.sign(np_s)
|
|
np_log_s = np.log(np.abs(np_s))
|
|
np_u = np.triu(np_u, k=1)
|
|
l_mask = np.tril(np.ones(w_shape, dtype=float), -1)
|
|
eye = np.eye(*w_shape, dtype=float)
|
|
|
|
self.register_buffer('p', torch.Tensor(np_p.astype(float)))
|
|
self.register_buffer('sign_s', torch.Tensor(np_sign_s.astype(float)))
|
|
self.l = nn.Parameter(torch.Tensor(np_l.astype(float)))
|
|
self.log_s = nn.Parameter(torch.Tensor(np_log_s.astype(float)))
|
|
self.u = nn.Parameter(torch.Tensor(np_u.astype(float)))
|
|
self.l_mask = torch.Tensor(l_mask)
|
|
self.eye = torch.Tensor(eye)
|
|
self.w_shape = w_shape
|
|
self.LU = LU_decomposed
|
|
self.weight = None
|
|
|
|
def get_weight(self, device, reverse):
|
|
w_shape = self.w_shape
|
|
self.p = self.p.to(device)
|
|
self.sign_s = self.sign_s.to(device)
|
|
self.l_mask = self.l_mask.to(device)
|
|
self.eye = self.eye.to(device)
|
|
l = self.l * self.l_mask + self.eye
|
|
u = self.u * self.l_mask.transpose(0, 1).contiguous() + torch.diag(self.sign_s * torch.exp(self.log_s))
|
|
dlogdet = self.log_s.sum()
|
|
if not reverse:
|
|
w = torch.matmul(self.p, torch.matmul(l, u))
|
|
else:
|
|
l = torch.inverse(l.double()).float()
|
|
u = torch.inverse(u.double()).float()
|
|
w = torch.matmul(u, torch.matmul(l, self.p.inverse()))
|
|
return w.view(w_shape[0], w_shape[1], 1), dlogdet
|
|
|
|
def forward(self, x, x_mask=None, reverse=False, **kwargs):
|
|
"""
|
|
log-det = log|abs(|W|)| * pixels
|
|
"""
|
|
b, c, t = x.size()
|
|
if x_mask is None:
|
|
x_len = torch.ones((b,), dtype=x.dtype, device=x.device) * t
|
|
else:
|
|
x_len = torch.sum(x_mask, [1, 2])
|
|
logdet = 0
|
|
if not reverse:
|
|
weight, dlogdet = self.get_weight(x.device, reverse)
|
|
z = F.conv1d(x, weight)
|
|
if logdet is not None:
|
|
logdet = logdet + dlogdet * x_len
|
|
return z, logdet
|
|
else:
|
|
if self.weight is None:
|
|
weight, dlogdet = self.get_weight(x.device, reverse)
|
|
else:
|
|
weight, dlogdet = self.weight, self.dlogdet
|
|
z = F.conv1d(x, weight)
|
|
if logdet is not None:
|
|
logdet = logdet - dlogdet * x_len
|
|
return z, logdet
|
|
|
|
def store_inverse(self):
|
|
self.weight, self.dlogdet = self.get_weight('cuda', reverse=True)
|
|
|
|
|
|
class Flip(nn.Module):
|
|
def forward(self, x, *args, reverse=False, **kwargs):
|
|
x = torch.flip(x, [1])
|
|
logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
|
|
return x, logdet
|
|
|
|
def store_inverse(self):
|
|
pass
|
|
|
|
|
|
class CouplingBlock(nn.Module): # 仿射耦合层
|
|
def __init__(self, in_channels, hidden_channels, kernel_size, dilation_rate, n_layers,
|
|
gin_channels=0, p_dropout=0, sigmoid_scale=False,
|
|
share_cond_layers=False, wn=None):
|
|
super().__init__()
|
|
self.in_channels = in_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.p_dropout = p_dropout
|
|
self.sigmoid_scale = sigmoid_scale
|
|
|
|
start = torch.nn.Conv1d(in_channels // 2, hidden_channels, 1)
|
|
start = torch.nn.utils.weight_norm(start)
|
|
self.start = start
|
|
# Initializing last layer to 0 makes the affine coupling layers
|
|
# do nothing at first. This helps with training stability
|
|
end = torch.nn.Conv1d(hidden_channels, in_channels, 1)
|
|
end.weight.data.zero_()
|
|
end.bias.data.zero_()
|
|
self.end = end
|
|
self.wn = WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels,
|
|
p_dropout, share_cond_layers)
|
|
if wn is not None:
|
|
self.wn.in_layers = wn.in_layers
|
|
self.wn.res_skip_layers = wn.res_skip_layers
|
|
|
|
def forward(self, x, x_mask=None, reverse=False, g=None, **kwargs):
|
|
if x_mask is None:
|
|
x_mask = 1
|
|
x_0, x_1 = x[:, :self.in_channels // 2], x[:, self.in_channels // 2:]
|
|
|
|
x = self.start(x_0) * x_mask
|
|
x = self.wn(x, x_mask, g)
|
|
out = self.end(x)
|
|
|
|
z_0 = x_0
|
|
m = out[:, :self.in_channels // 2, :]
|
|
logs = out[:, self.in_channels // 2:, :]
|
|
if self.sigmoid_scale:
|
|
logs = torch.log(1e-6 + torch.sigmoid(logs + 2))
|
|
if reverse:
|
|
z_1 = (x_1 - m) * torch.exp(-logs) * x_mask
|
|
logdet = torch.sum(-logs * x_mask, [1, 2])
|
|
else:
|
|
z_1 = (m + torch.exp(logs) * x_1) * x_mask
|
|
logdet = torch.sum(logs * x_mask, [1, 2])
|
|
z = torch.cat([z_0, z_1], 1)
|
|
return z, logdet
|
|
|
|
def store_inverse(self):
|
|
self.wn.remove_weight_norm()
|
|
|
|
|
|
class GlowFFTBlocks(FFTBlocks):
|
|
def __init__(self, hidden_size=128, gin_channels=256, num_layers=2, ffn_kernel_size=5,
|
|
dropout=None, num_heads=4, use_pos_embed=True, use_last_norm=True,
|
|
norm='ln', use_pos_embed_alpha=True):
|
|
super().__init__(hidden_size, num_layers, ffn_kernel_size, dropout, num_heads, use_pos_embed,
|
|
use_last_norm, norm, use_pos_embed_alpha)
|
|
self.inp_proj = nn.Conv1d(hidden_size + gin_channels, hidden_size, 1)
|
|
|
|
def forward(self, x, x_mask=None, g=None):
|
|
"""
|
|
:param x: [B, C_x, T]
|
|
:param x_mask: [B, 1, T]
|
|
:param g: [B, C_g, T]
|
|
:return: [B, C_x, T]
|
|
"""
|
|
if g is not None:
|
|
x = self.inp_proj(torch.cat([x, g], 1))
|
|
x = x.transpose(1, 2)
|
|
x = super(GlowFFTBlocks, self).forward(x, x_mask[:, 0] == 0)
|
|
x = x.transpose(1, 2)
|
|
return x
|
|
|
|
|
|
class TransformerCouplingBlock(nn.Module):
|
|
def __init__(self, in_channels, hidden_channels, n_layers,
|
|
gin_channels=0, p_dropout=0, sigmoid_scale=False):
|
|
super().__init__()
|
|
self.in_channels = in_channels
|
|
self.hidden_channels = hidden_channels
|
|
self.n_layers = n_layers
|
|
self.gin_channels = gin_channels
|
|
self.p_dropout = p_dropout
|
|
self.sigmoid_scale = sigmoid_scale
|
|
|
|
start = torch.nn.Conv1d(in_channels // 2, hidden_channels, 1)
|
|
self.start = start
|
|
# Initializing last layer to 0 makes the affine coupling layers
|
|
# do nothing at first. This helps with training stability
|
|
end = torch.nn.Conv1d(hidden_channels, in_channels, 1)
|
|
end.weight.data.zero_()
|
|
end.bias.data.zero_()
|
|
self.end = end
|
|
self.fft_blocks = GlowFFTBlocks(
|
|
hidden_size=hidden_channels,
|
|
ffn_kernel_size=3,
|
|
gin_channels=gin_channels,
|
|
num_layers=n_layers)
|
|
|
|
def forward(self, x, x_mask=None, reverse=False, g=None, **kwargs):
|
|
if x_mask is None:
|
|
x_mask = 1
|
|
x_0, x_1 = x[:, :self.in_channels // 2], x[:, self.in_channels // 2:]
|
|
|
|
x = self.start(x_0) * x_mask
|
|
x = self.fft_blocks(x, x_mask, g)
|
|
out = self.end(x)
|
|
|
|
z_0 = x_0
|
|
m = out[:, :self.in_channels // 2, :]
|
|
logs = out[:, self.in_channels // 2:, :]
|
|
if self.sigmoid_scale:
|
|
logs = torch.log(1e-6 + torch.sigmoid(logs + 2))
|
|
if reverse:
|
|
z_1 = (x_1 - m) * torch.exp(-logs) * x_mask
|
|
logdet = torch.sum(-logs * x_mask, [1, 2])
|
|
else:
|
|
z_1 = (m + torch.exp(logs) * x_1) * x_mask
|
|
logdet = torch.sum(logs * x_mask, [1, 2])
|
|
z = torch.cat([z_0, z_1], 1)
|
|
return z, logdet
|
|
|
|
def store_inverse(self):
|
|
pass
|
|
|
|
|
|
class FreqFFTCouplingBlock(nn.Module):
|
|
def __init__(self, in_channels, hidden_channels, n_layers,
|
|
gin_channels=0, p_dropout=0, sigmoid_scale=False):
|
|
super().__init__()
|
|
self.in_channels = in_channels
|
|
self.hidden_channels = hidden_channels
|
|
self.n_layers = n_layers
|
|
self.gin_channels = gin_channels
|
|
self.p_dropout = p_dropout
|
|
self.sigmoid_scale = sigmoid_scale
|
|
|
|
hs = hidden_channels
|
|
stride = 8
|
|
self.start = torch.nn.Conv2d(3, hs, kernel_size=stride * 2,
|
|
stride=stride, padding=stride // 2)
|
|
end = nn.ConvTranspose2d(hs, 2, kernel_size=stride, stride=stride)
|
|
end.weight.data.zero_()
|
|
end.bias.data.zero_()
|
|
self.end = nn.Sequential(
|
|
nn.Conv2d(hs * 3, hs, 3, 1, 1),
|
|
nn.ReLU(),
|
|
nn.GroupNorm(4, hs),
|
|
nn.Conv2d(hs, hs, 3, 1, 1),
|
|
end
|
|
)
|
|
self.fft_v = FFTBlocks(hidden_size=hs, ffn_kernel_size=1, num_layers=n_layers)
|
|
self.fft_h = nn.Sequential(
|
|
nn.Conv1d(hs, hs, 3, 1, 1),
|
|
nn.ReLU(),
|
|
nn.Conv1d(hs, hs, 3, 1, 1),
|
|
)
|
|
self.fft_g = nn.Sequential(
|
|
nn.Conv1d(
|
|
gin_channels - 160, hs, kernel_size=stride * 2, stride=stride, padding=stride // 2),
|
|
Permute(0, 2, 1),
|
|
FFTBlocks(hidden_size=hs, ffn_kernel_size=1, num_layers=n_layers),
|
|
Permute(0, 2, 1),
|
|
)
|
|
|
|
def forward(self, x, x_mask=None, reverse=False, g=None, **kwargs):
|
|
g_, _ = unsqueeze(g)
|
|
g_mel = g_[:, :80]
|
|
g_txt = g_[:, 80:]
|
|
g_mel, _ = squeeze(g_mel)
|
|
g_txt, _ = squeeze(g_txt) # [B, C, T]
|
|
|
|
if x_mask is None:
|
|
x_mask = 1
|
|
x_0, x_1 = x[:, :self.in_channels // 2], x[:, self.in_channels // 2:]
|
|
x = torch.stack([x_0, g_mel[:, :80], g_mel[:, 80:]], 1)
|
|
x = self.start(x) # [B, C, N_bins, T]
|
|
B, C, N_bins, T = x.shape
|
|
|
|
x_v = self.fft_v(x.permute(0, 3, 2, 1).reshape(B * T, N_bins, C))
|
|
x_v = x_v.reshape(B, T, N_bins, -1).permute(0, 3, 2, 1)
|
|
# x_v = x
|
|
|
|
x_h = self.fft_h(x.permute(0, 2, 1, 3).reshape(B * N_bins, C, T))
|
|
x_h = x_h.reshape(B, N_bins, -1, T).permute(0, 2, 1, 3)
|
|
# x_h = x
|
|
|
|
x_g = self.fft_g(g_txt)[:, :, None, :].repeat(1, 1, 10, 1)
|
|
x = torch.cat([x_v, x_h, x_g], 1)
|
|
out = self.end(x)
|
|
|
|
z_0 = x_0
|
|
m = out[:, 0]
|
|
logs = out[:, 1]
|
|
if self.sigmoid_scale:
|
|
logs = torch.log(1e-6 + torch.sigmoid(logs + 2))
|
|
if reverse:
|
|
z_1 = (x_1 - m) * torch.exp(-logs) * x_mask
|
|
logdet = torch.sum(-logs * x_mask, [1, 2])
|
|
else:
|
|
z_1 = (m + torch.exp(logs) * x_1) * x_mask
|
|
logdet = torch.sum(logs * x_mask, [1, 2])
|
|
z = torch.cat([z_0, z_1], 1)
|
|
return z, logdet
|
|
|
|
def store_inverse(self):
|
|
pass
|
|
|
|
|
|
class Glow(nn.Module):
|
|
def __init__(self,
|
|
in_channels,
|
|
hidden_channels,
|
|
kernel_size,
|
|
dilation_rate,
|
|
n_blocks,
|
|
n_layers,
|
|
p_dropout=0.,
|
|
n_split=4,
|
|
n_sqz=2,
|
|
sigmoid_scale=False,
|
|
gin_channels=0,
|
|
inv_conv_type='near',
|
|
share_cond_layers=False,
|
|
share_wn_layers=0,
|
|
):
|
|
super().__init__()
|
|
|
|
self.in_channels = in_channels
|
|
self.hidden_channels = hidden_channels
|
|
self.kernel_size = kernel_size
|
|
self.dilation_rate = dilation_rate
|
|
self.n_blocks = n_blocks
|
|
self.n_layers = n_layers
|
|
self.p_dropout = p_dropout
|
|
self.n_split = n_split
|
|
self.n_sqz = n_sqz
|
|
self.sigmoid_scale = sigmoid_scale
|
|
self.gin_channels = gin_channels
|
|
self.share_cond_layers = share_cond_layers
|
|
if gin_channels != 0 and share_cond_layers:
|
|
cond_layer = torch.nn.Conv1d(gin_channels * n_sqz, 2 * hidden_channels * n_layers, 1)
|
|
self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')
|
|
wn = None
|
|
self.flows = nn.ModuleList()
|
|
for b in range(n_blocks):
|
|
self.flows.append(ActNorm(channels=in_channels * n_sqz))
|
|
if inv_conv_type == 'near':
|
|
self.flows.append(InvConvNear(channels=in_channels * n_sqz, n_split=n_split, n_sqz=n_sqz))
|
|
if inv_conv_type == 'invconv':
|
|
self.flows.append(InvConv(channels=in_channels * n_sqz))
|
|
if share_wn_layers > 0:
|
|
if b % share_wn_layers == 0:
|
|
wn = WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels * n_sqz,
|
|
p_dropout, share_cond_layers)
|
|
self.flows.append(
|
|
CouplingBlock(
|
|
in_channels * n_sqz,
|
|
hidden_channels,
|
|
kernel_size=kernel_size,
|
|
dilation_rate=dilation_rate,
|
|
n_layers=n_layers,
|
|
gin_channels=gin_channels * n_sqz,
|
|
p_dropout=p_dropout,
|
|
sigmoid_scale=sigmoid_scale,
|
|
share_cond_layers=share_cond_layers,
|
|
wn=wn
|
|
))
|
|
|
|
def forward(self, x, x_mask=None, g=None, reverse=False, return_hiddens=False):
|
|
logdet_tot = 0
|
|
if not reverse:
|
|
flows = self.flows
|
|
else:
|
|
flows = reversed(self.flows)
|
|
if return_hiddens:
|
|
hs = []
|
|
if self.n_sqz > 1:
|
|
x, x_mask_ = squeeze(x, x_mask, self.n_sqz)
|
|
if g is not None:
|
|
g, _ = squeeze(g, x_mask, self.n_sqz)
|
|
x_mask = x_mask_
|
|
if self.share_cond_layers and g is not None:
|
|
g = self.cond_layer(g)
|
|
for f in flows:
|
|
x, logdet = f(x, x_mask, g=g, reverse=reverse)
|
|
if return_hiddens:
|
|
hs.append(x)
|
|
logdet_tot += logdet
|
|
if self.n_sqz > 1:
|
|
x, x_mask = unsqueeze(x, x_mask, self.n_sqz)
|
|
if return_hiddens:
|
|
return x, logdet_tot, hs
|
|
return x, logdet_tot
|
|
|
|
def store_inverse(self):
|
|
def remove_weight_norm(m):
|
|
try:
|
|
nn.utils.remove_weight_norm(m)
|
|
except ValueError: # this module didn't have weight norm
|
|
return
|
|
|
|
self.apply(remove_weight_norm)
|
|
for f in self.flows:
|
|
f.store_inverse()
|
|
|
|
|
|
class GlowV2(nn.Module):
|
|
def __init__(self,
|
|
in_channels=256,
|
|
hidden_channels=256,
|
|
kernel_size=3,
|
|
dilation_rate=1,
|
|
n_blocks=8,
|
|
n_layers=4,
|
|
p_dropout=0.,
|
|
n_split=4,
|
|
n_split_blocks=3,
|
|
sigmoid_scale=False,
|
|
gin_channels=0,
|
|
share_cond_layers=True):
|
|
super().__init__()
|
|
|
|
self.in_channels = in_channels
|
|
self.hidden_channels = hidden_channels
|
|
self.kernel_size = kernel_size
|
|
self.dilation_rate = dilation_rate
|
|
self.n_blocks = n_blocks
|
|
self.n_layers = n_layers
|
|
self.p_dropout = p_dropout
|
|
self.n_split = n_split
|
|
self.n_split_blocks = n_split_blocks
|
|
self.sigmoid_scale = sigmoid_scale
|
|
self.gin_channels = gin_channels
|
|
|
|
self.cond_layers = nn.ModuleList()
|
|
self.share_cond_layers = share_cond_layers
|
|
|
|
self.flows = nn.ModuleList()
|
|
in_channels = in_channels * 2
|
|
for l in range(n_split_blocks):
|
|
blocks = nn.ModuleList()
|
|
self.flows.append(blocks)
|
|
gin_channels = gin_channels * 2
|
|
if gin_channels != 0 and share_cond_layers:
|
|
cond_layer = torch.nn.Conv1d(gin_channels, 2 * hidden_channels * n_layers, 1)
|
|
self.cond_layers.append(torch.nn.utils.weight_norm(cond_layer, name='weight'))
|
|
for b in range(n_blocks):
|
|
blocks.append(ActNorm(channels=in_channels))
|
|
blocks.append(InvConvNear(channels=in_channels, n_split=n_split))
|
|
blocks.append(CouplingBlock(
|
|
in_channels,
|
|
hidden_channels,
|
|
kernel_size=kernel_size,
|
|
dilation_rate=dilation_rate,
|
|
n_layers=n_layers,
|
|
gin_channels=gin_channels,
|
|
p_dropout=p_dropout,
|
|
sigmoid_scale=sigmoid_scale,
|
|
share_cond_layers=share_cond_layers))
|
|
|
|
def forward(self, x=None, x_mask=None, g=None, reverse=False, concat_zs=True,
|
|
noise_scale=0.66, return_hiddens=False):
|
|
logdet_tot = 0
|
|
if not reverse:
|
|
flows = self.flows
|
|
assert x_mask is not None
|
|
zs = []
|
|
if return_hiddens:
|
|
hs = []
|
|
for i, blocks in enumerate(flows):
|
|
x, x_mask = squeeze(x, x_mask)
|
|
g_ = None
|
|
if g is not None:
|
|
g, _ = squeeze(g)
|
|
if self.share_cond_layers:
|
|
g_ = self.cond_layers[i](g)
|
|
else:
|
|
g_ = g
|
|
for layer in blocks:
|
|
x, logdet = layer(x, x_mask=x_mask, g=g_, reverse=reverse)
|
|
if return_hiddens:
|
|
hs.append(x)
|
|
logdet_tot += logdet
|
|
if i == self.n_split_blocks - 1:
|
|
zs.append(x)
|
|
else:
|
|
x, z = torch.chunk(x, 2, 1)
|
|
zs.append(z)
|
|
if concat_zs:
|
|
zs = [z.reshape(x.shape[0], -1) for z in zs]
|
|
zs = torch.cat(zs, 1) # [B, C*T]
|
|
if return_hiddens:
|
|
return zs, logdet_tot, hs
|
|
return zs, logdet_tot
|
|
else:
|
|
flows = reversed(self.flows)
|
|
if x is not None:
|
|
assert isinstance(x, list)
|
|
zs = x
|
|
else:
|
|
B, _, T = g.shape
|
|
zs = self.get_prior(B, T, g.device, noise_scale)
|
|
zs_ori = zs
|
|
if g is not None:
|
|
g_, g = g, []
|
|
for i in range(len(self.flows)):
|
|
g_, _ = squeeze(g_)
|
|
g.append(self.cond_layers[i](g_) if self.share_cond_layers else g_)
|
|
else:
|
|
g = [None for _ in range(len(self.flows))]
|
|
if x_mask is not None:
|
|
x_masks = []
|
|
for i in range(len(self.flows)):
|
|
x_mask, _ = squeeze(x_mask)
|
|
x_masks.append(x_mask)
|
|
else:
|
|
x_masks = [None for _ in range(len(self.flows))]
|
|
x_masks = x_masks[::-1]
|
|
g = g[::-1]
|
|
zs = zs[::-1]
|
|
x = None
|
|
for i, blocks in enumerate(flows):
|
|
x = zs[i] if x is None else torch.cat([x, zs[i]], 1)
|
|
for layer in reversed(blocks):
|
|
x, logdet = layer(x, x_masks=x_masks[i], g=g[i], reverse=reverse)
|
|
logdet_tot += logdet
|
|
x, _ = unsqueeze(x)
|
|
return x, logdet_tot, zs_ori
|
|
|
|
def store_inverse(self):
|
|
for f in self.modules():
|
|
if hasattr(f, 'store_inverse') and f != self:
|
|
f.store_inverse()
|
|
|
|
def remove_weight_norm(m):
|
|
try:
|
|
nn.utils.remove_weight_norm(m)
|
|
except ValueError: # this module didn't have weight norm
|
|
return
|
|
|
|
self.apply(remove_weight_norm)
|
|
|
|
def get_prior(self, B, T, device, noise_scale=0.66):
|
|
C = 80
|
|
zs = []
|
|
for i in range(len(self.flows)):
|
|
C, T = C, T // 2
|
|
if i == self.n_split_blocks - 1:
|
|
zs.append(torch.randn(B, C * 2, T).to(device) * noise_scale)
|
|
else:
|
|
zs.append(torch.randn(B, C, T).to(device) * noise_scale)
|
|
return zs
|
|
|
|
|
|
def squeeze(x, x_mask=None, n_sqz=2):
|
|
b, c, t = x.size()
|
|
|
|
t = (t // n_sqz) * n_sqz
|
|
x = x[:, :, :t]
|
|
x_sqz = x.view(b, c, t // n_sqz, n_sqz)
|
|
x_sqz = x_sqz.permute(0, 3, 1, 2).contiguous().view(b, c * n_sqz, t // n_sqz)
|
|
|
|
if x_mask is not None:
|
|
x_mask = x_mask[:, :, n_sqz - 1::n_sqz]
|
|
else:
|
|
x_mask = torch.ones(b, 1, t // n_sqz).to(device=x.device, dtype=x.dtype)
|
|
return x_sqz * x_mask, x_mask
|
|
|
|
|
|
def unsqueeze(x, x_mask=None, n_sqz=2):
|
|
b, c, t = x.size()
|
|
|
|
x_unsqz = x.view(b, n_sqz, c // n_sqz, t)
|
|
x_unsqz = x_unsqz.permute(0, 2, 3, 1).contiguous().view(b, c // n_sqz, t * n_sqz)
|
|
|
|
if x_mask is not None:
|
|
x_mask = x_mask.unsqueeze(-1).repeat(1, 1, 1, n_sqz).view(b, 1, t * n_sqz)
|
|
else:
|
|
x_mask = torch.ones(b, 1, t * n_sqz).to(device=x.device, dtype=x.dtype)
|
|
return x_unsqz * x_mask, x_mask
|