| from __future__ import annotations |
| import math |
| import copy |
| import torch |
| import numpy as np |
| from torch import nn |
| from torch.nn import functional as F |
| from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d |
| from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm |
|
|
|
|
| import math |
| import numpy as np |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| def init_weights(m, mean=0.0, std=0.01): |
| classname = m.__class__.__name__ |
| if classname.find("Conv") != -1: |
| m.weight.data.normal_(mean, std) |
|
|
|
|
| def get_padding(kernel_size, dilation=1): |
| return int((kernel_size*dilation - dilation)/2) |
|
|
|
|
| def convert_pad_shape(pad_shape): |
| l = pad_shape[::-1] |
| pad_shape = [item for sublist in l for item in sublist] |
| return pad_shape |
|
|
|
|
| def intersperse(lst, item): |
| result = [item] * (len(lst) * 2 + 1) |
| result[1::2] = lst |
| return result |
|
|
|
|
| def kl_divergence(m_p, logs_p, m_q, logs_q): |
| """KL(P||Q)""" |
| kl = (logs_q - logs_p) - 0.5 |
| kl += 0.5 * (torch.exp(2. * logs_p) + ((m_p - m_q)**2)) * torch.exp(-2. * logs_q) |
| return kl |
|
|
|
|
| def rand_gumbel(shape): |
| """Sample from the Gumbel distribution, protect from overflows.""" |
| uniform_samples = torch.rand(shape) * 0.99998 + 0.00001 |
| return -torch.log(-torch.log(uniform_samples)) |
|
|
|
|
| def rand_gumbel_like(x): |
| g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device) |
| return g |
|
|
|
|
| def slice_segments(x, ids_str, segment_size=4): |
| ret = torch.zeros_like(x[:, :, :segment_size]) |
| for i in range(x.size(0)): |
| idx_str = ids_str[i] |
| idx_end = idx_str + segment_size |
| ret[i] = x[i, :, idx_str:idx_end] |
| return ret |
|
|
|
|
| def maximum_path(neg_cent, mask): |
| device = neg_cent.device |
| dtype = neg_cent.dtype |
| neg_cent = neg_cent.detach().cpu().numpy().astype(np.float64) |
| mask = mask.detach().cpu().numpy().astype(np.bool_) |
| b, t_t, t_s = neg_cent.shape |
| path = np.zeros_like(neg_cent, dtype=np.float64) |
|
|
| for i in range(b): |
| t_y = int(mask[i, :, 0].sum()) |
| t_x = int(mask[i, 0, :].sum()) |
| if t_y == 0 or t_x == 0: |
| continue |
| value = neg_cent[i, :t_y, :t_x] |
| v = np.full((t_y, t_x), -1e9, dtype=np.float64) |
| v[0, 0] = value[0, 0] |
| for y in range(1, t_y): |
| v[y, 0] = v[y - 1, 0] + value[y, 0] |
| for x in range(1, t_x): |
| v[0, x] = -1e9 |
| for y in range(1, t_y): |
| for x in range(1, min(t_x, y + 1)): |
| v[y, x] = value[y, x] + max(v[y - 1, x], v[y - 1, x - 1]) |
| index = t_x - 1 |
| for y in range(t_y - 1, -1, -1): |
| path[i, y, index] = 1.0 |
| if index > 0 and (y == 0 or v[y - 1, index - 1] > v[y - 1, index]): |
| index -= 1 |
| return torch.from_numpy(path).to(device=device, dtype=dtype) |
|
|
|
|
| def rand_slice_segments(x, x_lengths=None, segment_size=4): |
| b, d, t = x.size() |
| if x_lengths is None: |
| x_lengths = t |
| ids_str_max = x_lengths - segment_size + 1 |
| ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long) |
| ret = slice_segments(x, ids_str, segment_size) |
| return ret, ids_str |
|
|
|
|
| def get_timing_signal_1d( |
| length, channels, min_timescale=1.0, max_timescale=1.0e4): |
| position = torch.arange(length, dtype=torch.float) |
| num_timescales = channels // 2 |
| log_timescale_increment = ( |
| math.log(float(max_timescale) / float(min_timescale)) / |
| (num_timescales - 1)) |
| inv_timescales = min_timescale * torch.exp( |
| torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment) |
| scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1) |
| signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0) |
| signal = F.pad(signal, [0, 0, 0, channels % 2]) |
| signal = signal.view(1, channels, length) |
| return signal |
|
|
|
|
| def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4): |
| b, channels, length = x.size() |
| signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale) |
| return x + signal.to(dtype=x.dtype, device=x.device) |
|
|
|
|
| def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1): |
| b, channels, length = x.size() |
| signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale) |
| return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis) |
|
|
|
|
| def subsequent_mask(length): |
| mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0) |
| return mask |
|
|
|
|
| @torch.jit.script |
| def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels): |
| n_channels_int = n_channels[0] |
| in_act = input_a + input_b |
| t_act = torch.tanh(in_act[:, :n_channels_int, :]) |
| s_act = torch.sigmoid(in_act[:, n_channels_int:, :]) |
| acts = t_act * s_act |
| return acts |
|
|
|
|
| def convert_pad_shape(pad_shape): |
| l = pad_shape[::-1] |
| pad_shape = [item for sublist in l for item in sublist] |
| return pad_shape |
|
|
|
|
| def shift_1d(x): |
| x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1] |
| return x |
|
|
|
|
| def sequence_mask(length, max_length=None): |
| if max_length is None: |
| max_length = length.max() |
| x = torch.arange(max_length, dtype=length.dtype, device=length.device) |
| return x.unsqueeze(0) < length.unsqueeze(1) |
|
|
|
|
| def generate_path(duration, mask): |
| """ |
| duration: [b, 1, t_x] |
| mask: [b, 1, t_y, t_x] |
| """ |
| device = duration.device |
| |
| b, _, t_y, t_x = mask.shape |
| cum_duration = torch.cumsum(duration, -1) |
| |
| cum_duration_flat = cum_duration.view(b * t_x) |
| path = sequence_mask(cum_duration_flat, t_y).to(mask.dtype) |
| path = path.view(b, t_x, t_y) |
| path = path - F.pad(path, convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1] |
| path = path.unsqueeze(1).transpose(2,3) * mask |
| return path |
|
|
|
|
| def clip_grad_value_(parameters, clip_value, norm_type=2): |
| if isinstance(parameters, torch.Tensor): |
| parameters = [parameters] |
| parameters = list(filter(lambda p: p.grad is not None, parameters)) |
| norm_type = float(norm_type) |
| if clip_value is not None: |
| clip_value = float(clip_value) |
|
|
| total_norm = 0 |
| for p in parameters: |
| param_norm = p.grad.data.norm(norm_type) |
| total_norm += param_norm.item() ** norm_type |
| if clip_value is not None: |
| p.grad.data.clamp_(min=-clip_value, max=clip_value) |
| total_norm = total_norm ** (1. / norm_type) |
| return total_norm |
|
|
| """Lightweight alias-free waveform blocks derived from NVIDIA BigVGAN. |
| |
| BigVGAN and alias-free-torch are MIT/Apache-2.0 licensed. The implementation |
| is kept local so Inflect can train without BigVGAN's optional CUDA extension. |
| """ |
|
|
| import math |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| from torch.nn.utils import remove_weight_norm, weight_norm |
|
|
|
|
|
|
| def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int): |
| even = kernel_size % 2 == 0 |
| half_size = kernel_size // 2 |
| delta_f = 4 * half_width |
| attenuation = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 |
| if attenuation > 50.0: |
| beta = 0.1102 * (attenuation - 8.7) |
| elif attenuation >= 21.0: |
| beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0) |
| else: |
| beta = 0.0 |
| window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) |
| if even: |
| time = torch.arange(-half_size, half_size) + 0.5 |
| else: |
| time = torch.arange(kernel_size) - half_size |
| values = 2 * cutoff * window * torch.sinc(2 * cutoff * time) |
| values /= values.sum() |
| return values.view(1, 1, kernel_size) |
|
|
|
|
| class UpSample1d(nn.Module): |
| def __init__(self, ratio=2, kernel_size=12): |
| super().__init__() |
| self.ratio = ratio |
| self.stride = ratio |
| self.kernel_size = kernel_size |
| self.pad = kernel_size // ratio - 1 |
| self.pad_left = self.pad * ratio + (kernel_size - ratio) // 2 |
| self.pad_right = self.pad * ratio + (kernel_size - ratio + 1) // 2 |
| self.register_buffer( |
| "filter", |
| kaiser_sinc_filter1d(0.5 / ratio, 0.6 / ratio, kernel_size)) |
|
|
| def forward(self, x): |
| channels = x.shape[1] |
| x = F.pad(x, (self.pad, self.pad), mode="replicate") |
| x = self.ratio * F.conv_transpose1d( |
| x, self.filter.expand(channels, -1, -1), |
| stride=self.stride, groups=channels) |
| return x[..., self.pad_left:-self.pad_right] |
|
|
|
|
| class DownSample1d(nn.Module): |
| def __init__(self, ratio=2, kernel_size=12): |
| super().__init__() |
| self.ratio = ratio |
| self.kernel_size = kernel_size |
| self.pad_left = kernel_size // 2 - int(kernel_size % 2 == 0) |
| self.pad_right = kernel_size // 2 |
| self.register_buffer( |
| "filter", |
| kaiser_sinc_filter1d(0.5 / ratio, 0.6 / ratio, kernel_size)) |
|
|
| def forward(self, x): |
| channels = x.shape[1] |
| x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate") |
| return F.conv1d( |
| x, self.filter.expand(channels, -1, -1), |
| stride=self.ratio, groups=channels) |
|
|
|
|
| class SnakeBeta(nn.Module): |
| def __init__(self, channels: int, logscale: bool = True): |
| super().__init__() |
| initial = torch.zeros(channels) if logscale else torch.ones(channels) |
| self.alpha = nn.Parameter(initial.clone()) |
| self.beta = nn.Parameter(initial.clone()) |
| self.logscale = logscale |
|
|
| def forward(self, x): |
| alpha = self.alpha.view(1, -1, 1) |
| beta = self.beta.view(1, -1, 1) |
| if self.logscale: |
| alpha = alpha.exp() |
| beta = beta.exp() |
| return x + torch.sin(x * alpha).square() / (beta + 1e-9) |
|
|
|
|
| class AliasFreeActivation1d(nn.Module): |
| def __init__(self, activation: nn.Module): |
| super().__init__() |
| self.upsample = UpSample1d() |
| self.act = activation |
| self.downsample = DownSample1d() |
|
|
| def forward(self, x): |
| return self.downsample(self.act(self.upsample(x))) |
|
|
|
|
| class AliasFreeResBlock1(nn.Module): |
| """Shape-compatible VITS ResBlock1 with filtered SnakeBeta activations.""" |
|
|
| def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), logscale=True): |
| super().__init__() |
| self.convs1 = nn.ModuleList([ |
| weight_norm(nn.Conv1d( |
| channels, channels, kernel_size, 1, |
| dilation=d, padding=get_padding(kernel_size, d))) |
| for d in dilation |
| ]) |
| self.convs2 = nn.ModuleList([ |
| weight_norm(nn.Conv1d( |
| channels, channels, kernel_size, 1, |
| dilation=1, padding=get_padding(kernel_size, 1))) |
| for _ in dilation |
| ]) |
| self.convs1.apply(init_weights) |
| self.convs2.apply(init_weights) |
| self.activations = nn.ModuleList([ |
| AliasFreeActivation1d(SnakeBeta(channels, logscale=logscale)) |
| for _ in range(2 * len(dilation)) |
| ]) |
|
|
| def forward(self, x, x_mask=None): |
| first = self.activations[::2] |
| second = self.activations[1::2] |
| for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, first, second): |
| residual = conv2(act2(conv1(act1(x)))) |
| x = x + residual |
| return x |
|
|
| def remove_weight_norm(self): |
| for layer in self.convs1: |
| remove_weight_norm(layer) |
| for layer in self.convs2: |
| remove_weight_norm(layer) |
|
|
| import copy |
| import math |
| import numpy as np |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| |
|
|
| class Encoder(nn.Module): |
| def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., window_size=4, **kwargs): |
| super().__init__() |
| self.hidden_channels = hidden_channels |
| self.filter_channels = filter_channels |
| self.n_heads = n_heads |
| self.n_layers = n_layers |
| self.kernel_size = kernel_size |
| self.p_dropout = p_dropout |
| self.window_size = window_size |
|
|
| self.drop = nn.Dropout(p_dropout) |
| self.attn_layers = nn.ModuleList() |
| self.norm_layers_1 = nn.ModuleList() |
| self.ffn_layers = nn.ModuleList() |
| self.norm_layers_2 = nn.ModuleList() |
| for i in range(self.n_layers): |
| self.attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, window_size=window_size)) |
| self.norm_layers_1.append(LayerNorm(hidden_channels)) |
| self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout)) |
| self.norm_layers_2.append(LayerNorm(hidden_channels)) |
|
|
| def forward(self, x, x_mask): |
| attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1) |
| x = x * x_mask |
| for i in range(self.n_layers): |
| y = self.attn_layers[i](x, x, attn_mask) |
| y = self.drop(y) |
| x = self.norm_layers_1[i](x + y) |
|
|
| y = self.ffn_layers[i](x, x_mask) |
| y = self.drop(y) |
| x = self.norm_layers_2[i](x + y) |
| x = x * x_mask |
| return x |
|
|
|
|
| class Decoder(nn.Module): |
| def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., proximal_bias=False, proximal_init=True, **kwargs): |
| super().__init__() |
| self.hidden_channels = hidden_channels |
| self.filter_channels = filter_channels |
| self.n_heads = n_heads |
| self.n_layers = n_layers |
| self.kernel_size = kernel_size |
| self.p_dropout = p_dropout |
| self.proximal_bias = proximal_bias |
| self.proximal_init = proximal_init |
|
|
| self.drop = nn.Dropout(p_dropout) |
| self.self_attn_layers = nn.ModuleList() |
| self.norm_layers_0 = nn.ModuleList() |
| self.encdec_attn_layers = nn.ModuleList() |
| self.norm_layers_1 = nn.ModuleList() |
| self.ffn_layers = nn.ModuleList() |
| self.norm_layers_2 = nn.ModuleList() |
| for i in range(self.n_layers): |
| self.self_attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, proximal_bias=proximal_bias, proximal_init=proximal_init)) |
| self.norm_layers_0.append(LayerNorm(hidden_channels)) |
| self.encdec_attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout)) |
| self.norm_layers_1.append(LayerNorm(hidden_channels)) |
| self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout, causal=True)) |
| self.norm_layers_2.append(LayerNorm(hidden_channels)) |
|
|
| def forward(self, x, x_mask, h, h_mask): |
| """ |
| x: decoder input |
| h: encoder output |
| """ |
| self_attn_mask = subsequent_mask(x_mask.size(2)).to(device=x.device, dtype=x.dtype) |
| encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1) |
| x = x * x_mask |
| for i in range(self.n_layers): |
| y = self.self_attn_layers[i](x, x, self_attn_mask) |
| y = self.drop(y) |
| x = self.norm_layers_0[i](x + y) |
|
|
| y = self.encdec_attn_layers[i](x, h, encdec_attn_mask) |
| y = self.drop(y) |
| x = self.norm_layers_1[i](x + y) |
| |
| y = self.ffn_layers[i](x, x_mask) |
| y = self.drop(y) |
| x = self.norm_layers_2[i](x + y) |
| x = x * x_mask |
| return x |
|
|
|
|
| class MultiHeadAttention(nn.Module): |
| def __init__(self, channels, out_channels, n_heads, p_dropout=0., window_size=None, heads_share=True, block_length=None, proximal_bias=False, proximal_init=False): |
| super().__init__() |
| assert channels % n_heads == 0 |
|
|
| self.channels = channels |
| self.out_channels = out_channels |
| self.n_heads = n_heads |
| self.p_dropout = p_dropout |
| self.window_size = window_size |
| self.heads_share = heads_share |
| self.block_length = block_length |
| self.proximal_bias = proximal_bias |
| self.proximal_init = proximal_init |
| self.attn = None |
|
|
| self.k_channels = channels // n_heads |
| self.conv_q = nn.Conv1d(channels, channels, 1) |
| self.conv_k = nn.Conv1d(channels, channels, 1) |
| self.conv_v = nn.Conv1d(channels, channels, 1) |
| self.conv_o = nn.Conv1d(channels, out_channels, 1) |
| self.drop = nn.Dropout(p_dropout) |
|
|
| if window_size is not None: |
| n_heads_rel = 1 if heads_share else n_heads |
| rel_stddev = self.k_channels**-0.5 |
| self.emb_rel_k = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev) |
| self.emb_rel_v = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev) |
|
|
| nn.init.xavier_uniform_(self.conv_q.weight) |
| nn.init.xavier_uniform_(self.conv_k.weight) |
| nn.init.xavier_uniform_(self.conv_v.weight) |
| if proximal_init: |
| with torch.no_grad(): |
| self.conv_k.weight.copy_(self.conv_q.weight) |
| self.conv_k.bias.copy_(self.conv_q.bias) |
| |
| def forward(self, x, c, attn_mask=None): |
| q = self.conv_q(x) |
| k = self.conv_k(c) |
| v = self.conv_v(c) |
| |
| x, self.attn = self.attention(q, k, v, mask=attn_mask) |
|
|
| x = self.conv_o(x) |
| return x |
|
|
| def attention(self, query, key, value, mask=None): |
| |
| b, d, t_s, t_t = (*key.size(), query.size(2)) |
| query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3) |
| key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3) |
| value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3) |
|
|
| scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1)) |
| if self.window_size is not None: |
| assert t_s == t_t, "Relative attention is only available for self-attention." |
| key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s) |
| rel_logits = self._matmul_with_relative_keys(query /math.sqrt(self.k_channels), key_relative_embeddings) |
| scores_local = self._relative_position_to_absolute_position(rel_logits) |
| scores = scores + scores_local |
| if self.proximal_bias: |
| assert t_s == t_t, "Proximal bias is only available for self-attention." |
| scores = scores + self._attention_bias_proximal(t_s).to(device=scores.device, dtype=scores.dtype) |
| if mask is not None: |
| scores = scores.masked_fill(mask == 0, -1e4) |
| if self.block_length is not None: |
| assert t_s == t_t, "Local attention is only available for self-attention." |
| block_mask = torch.ones_like(scores).triu(-self.block_length).tril(self.block_length) |
| scores = scores.masked_fill(block_mask == 0, -1e4) |
| p_attn = F.softmax(scores, dim=-1) |
| p_attn = self.drop(p_attn) |
| output = torch.matmul(p_attn, value) |
| if self.window_size is not None: |
| relative_weights = self._absolute_position_to_relative_position(p_attn) |
| value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s) |
| output = output + self._matmul_with_relative_values(relative_weights, value_relative_embeddings) |
| output = output.transpose(2, 3).contiguous().view(b, d, t_t) |
| return output, p_attn |
|
|
| def _matmul_with_relative_values(self, x, y): |
| """ |
| x: [b, h, l, m] |
| y: [h or 1, m, d] |
| ret: [b, h, l, d] |
| """ |
| ret = torch.matmul(x, y.unsqueeze(0)) |
| return ret |
|
|
| def _matmul_with_relative_keys(self, x, y): |
| """ |
| x: [b, h, l, d] |
| y: [h or 1, m, d] |
| ret: [b, h, l, m] |
| """ |
| ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1)) |
| return ret |
|
|
| def _get_relative_embeddings(self, relative_embeddings, length): |
| max_relative_position = 2 * self.window_size + 1 |
| |
| pad_length = max(length - (self.window_size + 1), 0) |
| slice_start_position = max((self.window_size + 1) - length, 0) |
| slice_end_position = slice_start_position + 2 * length - 1 |
| if pad_length > 0: |
| padded_relative_embeddings = F.pad( |
| relative_embeddings, |
| convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]])) |
| else: |
| padded_relative_embeddings = relative_embeddings |
| used_relative_embeddings = padded_relative_embeddings[:,slice_start_position:slice_end_position] |
| return used_relative_embeddings |
|
|
| def _relative_position_to_absolute_position(self, x): |
| """ |
| x: [b, h, l, 2*l-1] |
| ret: [b, h, l, l] |
| """ |
| batch, heads, length, _ = x.size() |
| |
| x = F.pad(x, convert_pad_shape([[0,0],[0,0],[0,0],[0,1]])) |
|
|
| |
| x_flat = x.view([batch, heads, length * 2 * length]) |
| x_flat = F.pad(x_flat, convert_pad_shape([[0,0],[0,0],[0,length-1]])) |
|
|
| |
| x_final = x_flat.view([batch, heads, length+1, 2*length-1])[:, :, :length, length-1:] |
| return x_final |
|
|
| def _absolute_position_to_relative_position(self, x): |
| """ |
| x: [b, h, l, l] |
| ret: [b, h, l, 2*l-1] |
| """ |
| batch, heads, length, _ = x.size() |
| |
| x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length-1]])) |
| x_flat = x.view([batch, heads, length**2 + length*(length -1)]) |
| |
| x_flat = F.pad(x_flat, convert_pad_shape([[0, 0], [0, 0], [length, 0]])) |
| x_final = x_flat.view([batch, heads, length, 2*length])[:,:,:,1:] |
| return x_final |
|
|
| def _attention_bias_proximal(self, length): |
| """Bias for self-attention to encourage attention to close positions. |
| Args: |
| length: an integer scalar. |
| Returns: |
| a Tensor with shape [1, 1, length, length] |
| """ |
| r = torch.arange(length, dtype=torch.float32) |
| diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1) |
| return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0) |
|
|
|
|
| class FFN(nn.Module): |
| def __init__(self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0., activation=None, causal=False): |
| super().__init__() |
| self.in_channels = in_channels |
| self.out_channels = out_channels |
| self.filter_channels = filter_channels |
| self.kernel_size = kernel_size |
| self.p_dropout = p_dropout |
| self.activation = activation |
| self.causal = causal |
|
|
| if causal: |
| self.padding = self._causal_padding |
| else: |
| self.padding = self._same_padding |
|
|
| self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size) |
| self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size) |
| self.drop = nn.Dropout(p_dropout) |
|
|
| def forward(self, x, x_mask): |
| x = self.conv_1(self.padding(x * x_mask)) |
| if self.activation == "gelu": |
| x = x * torch.sigmoid(1.702 * x) |
| else: |
| x = torch.relu(x) |
| x = self.drop(x) |
| x = self.conv_2(self.padding(x * x_mask)) |
| return x * x_mask |
| |
| def _causal_padding(self, x): |
| if self.kernel_size == 1: |
| return x |
| pad_l = self.kernel_size - 1 |
| pad_r = 0 |
| padding = [[0, 0], [0, 0], [pad_l, pad_r]] |
| x = F.pad(x, convert_pad_shape(padding)) |
| return x |
|
|
| def _same_padding(self, x): |
| if self.kernel_size == 1: |
| return x |
| pad_l = (self.kernel_size - 1) // 2 |
| pad_r = self.kernel_size // 2 |
| padding = [[0, 0], [0, 0], [pad_l, pad_r]] |
| x = F.pad(x, convert_pad_shape(padding)) |
| return x |
|
|
| import copy |
| import math |
| import numpy as np |
| import scipy |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d |
| from torch.nn.utils import weight_norm, remove_weight_norm |
|
|
|
|
|
|
| LRELU_SLOPE = 0.1 |
|
|
|
|
| class LayerNorm(nn.Module): |
| def __init__(self, channels, eps=1e-5): |
| 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): |
| x = x.transpose(1, -1) |
| x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps) |
| return x.transpose(1, -1) |
|
|
| |
| 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 DDSConv(nn.Module): |
| """ |
| Dialted and Depth-Separable Convolution |
| """ |
| def __init__(self, channels, kernel_size, n_layers, p_dropout=0.): |
| super().__init__() |
| self.channels = channels |
| self.kernel_size = kernel_size |
| self.n_layers = n_layers |
| self.p_dropout = p_dropout |
|
|
| self.drop = nn.Dropout(p_dropout) |
| self.convs_sep = nn.ModuleList() |
| self.convs_1x1 = nn.ModuleList() |
| self.norms_1 = nn.ModuleList() |
| self.norms_2 = nn.ModuleList() |
| for i in range(n_layers): |
| dilation = kernel_size ** i |
| padding = (kernel_size * dilation - dilation) // 2 |
| self.convs_sep.append(nn.Conv1d(channels, channels, kernel_size, |
| groups=channels, dilation=dilation, padding=padding |
| )) |
| self.convs_1x1.append(nn.Conv1d(channels, channels, 1)) |
| self.norms_1.append(LayerNorm(channels)) |
| self.norms_2.append(LayerNorm(channels)) |
|
|
| def forward(self, x, x_mask, g=None): |
| if g is not None: |
| x = x + g |
| for i in range(self.n_layers): |
| y = self.convs_sep[i](x * x_mask) |
| y = self.norms_1[i](y) |
| y = F.gelu(y) |
| y = self.convs_1x1[i](y) |
| y = self.norms_2[i](y) |
| y = F.gelu(y) |
| y = self.drop(y) |
| x = x + y |
| return x * x_mask |
|
|
|
|
| class WN(torch.nn.Module): |
| def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0): |
| super(WN, self).__init__() |
| assert(kernel_size % 2 == 1) |
| 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.in_layers = torch.nn.ModuleList() |
| self.res_skip_layers = torch.nn.ModuleList() |
| self.drop = nn.Dropout(p_dropout) |
|
|
| if gin_channels != 0: |
| cond_layer = torch.nn.Conv1d(gin_channels, 2*hidden_channels*n_layers, 1) |
| self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight') |
|
|
| for i in range(n_layers): |
| dilation = dilation_rate ** i |
| padding = int((kernel_size * dilation - dilation) / 2) |
| in_layer = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, kernel_size, |
| dilation=dilation, padding=padding) |
| in_layer = torch.nn.utils.weight_norm(in_layer, name='weight') |
| self.in_layers.append(in_layer) |
|
|
| |
| if i < n_layers - 1: |
| res_skip_channels = 2 * hidden_channels |
| else: |
| res_skip_channels = hidden_channels |
|
|
| res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1) |
| res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight') |
| self.res_skip_layers.append(res_skip_layer) |
|
|
| def forward(self, x, x_mask, g=None, **kwargs): |
| output = torch.zeros_like(x) |
| n_channels_tensor = torch.IntTensor([self.hidden_channels]) |
|
|
| if g is not None: |
| g = self.cond_layer(g) |
|
|
| for i in range(self.n_layers): |
| x_in = self.in_layers[i](x) |
| if g is not None: |
| cond_offset = i * 2 * self.hidden_channels |
| g_l = g[:,cond_offset:cond_offset+2*self.hidden_channels,:] |
| else: |
| g_l = torch.zeros_like(x_in) |
|
|
| acts = fused_add_tanh_sigmoid_multiply( |
| x_in, |
| g_l, |
| n_channels_tensor) |
| acts = self.drop(acts) |
|
|
| res_skip_acts = self.res_skip_layers[i](acts) |
| if i < self.n_layers - 1: |
| res_acts = res_skip_acts[:,:self.hidden_channels,:] |
| x = (x + res_acts) * x_mask |
| output = output + res_skip_acts[:,self.hidden_channels:,:] |
| else: |
| output = output + res_skip_acts |
| return output * x_mask |
|
|
| def remove_weight_norm(self): |
| if self.gin_channels != 0: |
| torch.nn.utils.remove_weight_norm(self.cond_layer) |
| for l in self.in_layers: |
| torch.nn.utils.remove_weight_norm(l) |
| for l in self.res_skip_layers: |
| torch.nn.utils.remove_weight_norm(l) |
|
|
|
|
| class ResBlock1(torch.nn.Module): |
| def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)): |
| super(ResBlock1, 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) |
|
|
| def forward(self, x, x_mask=None): |
| for c1, c2 in zip(self.convs1, self.convs2): |
| xt = F.leaky_relu(x, LRELU_SLOPE) |
| if x_mask is not None: |
| xt = xt * x_mask |
| xt = c1(xt) |
| xt = F.leaky_relu(xt, LRELU_SLOPE) |
| if x_mask is not None: |
| xt = xt * x_mask |
| xt = c2(xt) |
| x = xt + x |
| if x_mask is not None: |
| x = x * x_mask |
| return x |
|
|
| def remove_weight_norm(self): |
| for l in self.convs1: |
| remove_weight_norm(l) |
| for l in self.convs2: |
| remove_weight_norm(l) |
|
|
|
|
| class ResBlock2(torch.nn.Module): |
| def __init__(self, channels, kernel_size=3, dilation=(1, 3)): |
| super(ResBlock2, self).__init__() |
| self.convs = 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]))) |
| ]) |
| self.convs.apply(init_weights) |
|
|
| def forward(self, x, x_mask=None): |
| for c in self.convs: |
| xt = F.leaky_relu(x, LRELU_SLOPE) |
| if x_mask is not None: |
| xt = xt * x_mask |
| xt = c(xt) |
| x = xt + x |
| if x_mask is not None: |
| x = x * x_mask |
| return x |
|
|
| def remove_weight_norm(self): |
| for l in self.convs: |
| remove_weight_norm(l) |
|
|
|
|
| class Log(nn.Module): |
| def forward(self, x, x_mask, reverse=False, **kwargs): |
| if not reverse: |
| y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask |
| logdet = torch.sum(-y, [1, 2]) |
| return y, logdet |
| else: |
| x = torch.exp(x) * x_mask |
| return x |
| |
|
|
| class Flip(nn.Module): |
| def forward(self, x, *args, reverse=False, **kwargs): |
| x = torch.flip(x, [1]) |
| if not reverse: |
| logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device) |
| return x, logdet |
| else: |
| return x |
|
|
|
|
| class ElementwiseAffine(nn.Module): |
| def __init__(self, channels): |
| super().__init__() |
| self.channels = channels |
| self.m = nn.Parameter(torch.zeros(channels,1)) |
| self.logs = nn.Parameter(torch.zeros(channels,1)) |
|
|
| def forward(self, x, x_mask, reverse=False, **kwargs): |
| if not reverse: |
| y = self.m + torch.exp(self.logs) * x |
| y = y * x_mask |
| logdet = torch.sum(self.logs * x_mask, [1,2]) |
| return y, logdet |
| else: |
| x = (x - self.m) * torch.exp(-self.logs) * x_mask |
| return x |
|
|
|
|
| class ResidualCouplingLayer(nn.Module): |
| def __init__(self, |
| channels, |
| hidden_channels, |
| kernel_size, |
| dilation_rate, |
| n_layers, |
| p_dropout=0, |
| gin_channels=0, |
| mean_only=False): |
| assert channels % 2 == 0, "channels should be divisible by 2" |
| super().__init__() |
| self.channels = channels |
| self.hidden_channels = hidden_channels |
| self.kernel_size = kernel_size |
| self.dilation_rate = dilation_rate |
| self.n_layers = n_layers |
| self.half_channels = channels // 2 |
| self.mean_only = mean_only |
|
|
| self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1) |
| self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=p_dropout, gin_channels=gin_channels) |
| self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1) |
| self.post.weight.data.zero_() |
| self.post.bias.data.zero_() |
|
|
| def forward(self, x, x_mask, g=None, reverse=False): |
| x0, x1 = torch.split(x, [self.half_channels]*2, 1) |
| h = self.pre(x0) * x_mask |
| h = self.enc(h, x_mask, g=g) |
| stats = self.post(h) * x_mask |
| if not self.mean_only: |
| m, logs = torch.split(stats, [self.half_channels]*2, 1) |
| else: |
| m = stats |
| logs = torch.zeros_like(m) |
|
|
| if not reverse: |
| x1 = m + x1 * torch.exp(logs) * x_mask |
| x = torch.cat([x0, x1], 1) |
| logdet = torch.sum(logs, [1,2]) |
| return x, logdet |
| else: |
| x1 = (x1 - m) * torch.exp(-logs) * x_mask |
| x = torch.cat([x0, x1], 1) |
| return x |
|
|
|
|
| class ConvFlow(nn.Module): |
| def __init__(self, in_channels, filter_channels, kernel_size, n_layers, num_bins=10, tail_bound=5.0): |
| super().__init__() |
| self.in_channels = in_channels |
| self.filter_channels = filter_channels |
| self.kernel_size = kernel_size |
| self.n_layers = n_layers |
| self.num_bins = num_bins |
| self.tail_bound = tail_bound |
| self.half_channels = in_channels // 2 |
|
|
| self.pre = nn.Conv1d(self.half_channels, filter_channels, 1) |
| self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.) |
| self.proj = nn.Conv1d(filter_channels, self.half_channels * (num_bins * 3 - 1), 1) |
| self.proj.weight.data.zero_() |
| self.proj.bias.data.zero_() |
|
|
| def forward(self, x, x_mask, g=None, reverse=False): |
| x0, x1 = torch.split(x, [self.half_channels]*2, 1) |
| h = self.pre(x0) |
| h = self.convs(h, x_mask, g=g) |
| h = self.proj(h) * x_mask |
|
|
| b, c, t = x0.shape |
| h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2) |
|
|
| unnormalized_widths = h[..., :self.num_bins] / math.sqrt(self.filter_channels) |
| unnormalized_heights = h[..., self.num_bins:2*self.num_bins] / math.sqrt(self.filter_channels) |
| unnormalized_derivatives = h[..., 2 * self.num_bins:] |
|
|
| x1, logabsdet = piecewise_rational_quadratic_transform(x1, |
| unnormalized_widths, |
| unnormalized_heights, |
| unnormalized_derivatives, |
| inverse=reverse, |
| tails='linear', |
| tail_bound=self.tail_bound |
| ) |
|
|
| x = torch.cat([x0, x1], 1) * x_mask |
| logdet = torch.sum(logabsdet * x_mask, [1,2]) |
| if not reverse: |
| return x, logdet |
| else: |
| return x |
|
|
| import torch |
| from torch.nn import functional as F |
|
|
| import numpy as np |
|
|
|
|
| DEFAULT_MIN_BIN_WIDTH = 1e-3 |
| DEFAULT_MIN_BIN_HEIGHT = 1e-3 |
| DEFAULT_MIN_DERIVATIVE = 1e-3 |
|
|
|
|
| def piecewise_rational_quadratic_transform(inputs, |
| unnormalized_widths, |
| unnormalized_heights, |
| unnormalized_derivatives, |
| inverse=False, |
| tails=None, |
| tail_bound=1., |
| min_bin_width=DEFAULT_MIN_BIN_WIDTH, |
| min_bin_height=DEFAULT_MIN_BIN_HEIGHT, |
| min_derivative=DEFAULT_MIN_DERIVATIVE): |
|
|
| if tails is None: |
| spline_fn = rational_quadratic_spline |
| spline_kwargs = {} |
| else: |
| spline_fn = unconstrained_rational_quadratic_spline |
| spline_kwargs = { |
| 'tails': tails, |
| 'tail_bound': tail_bound |
| } |
|
|
| outputs, logabsdet = spline_fn( |
| inputs=inputs, |
| unnormalized_widths=unnormalized_widths, |
| unnormalized_heights=unnormalized_heights, |
| unnormalized_derivatives=unnormalized_derivatives, |
| inverse=inverse, |
| min_bin_width=min_bin_width, |
| min_bin_height=min_bin_height, |
| min_derivative=min_derivative, |
| **spline_kwargs |
| ) |
| return outputs, logabsdet |
|
|
|
|
| def searchsorted(bin_locations, inputs, eps=1e-6): |
| bin_locations[..., -1] += eps |
| return torch.sum( |
| inputs[..., None] >= bin_locations, |
| dim=-1 |
| ) - 1 |
|
|
|
|
| def unconstrained_rational_quadratic_spline(inputs, |
| unnormalized_widths, |
| unnormalized_heights, |
| unnormalized_derivatives, |
| inverse=False, |
| tails='linear', |
| tail_bound=1., |
| min_bin_width=DEFAULT_MIN_BIN_WIDTH, |
| min_bin_height=DEFAULT_MIN_BIN_HEIGHT, |
| min_derivative=DEFAULT_MIN_DERIVATIVE): |
| inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound) |
| outside_interval_mask = ~inside_interval_mask |
|
|
| outputs = torch.zeros_like(inputs) |
| logabsdet = torch.zeros_like(inputs) |
|
|
| if tails == 'linear': |
| unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1)) |
| constant = np.log(np.exp(1 - min_derivative) - 1) |
| unnormalized_derivatives[..., 0] = constant |
| unnormalized_derivatives[..., -1] = constant |
|
|
| outputs[outside_interval_mask] = inputs[outside_interval_mask] |
| logabsdet[outside_interval_mask] = 0 |
| else: |
| raise RuntimeError('{} tails are not implemented.'.format(tails)) |
|
|
| outputs[inside_interval_mask], logabsdet[inside_interval_mask] = rational_quadratic_spline( |
| inputs=inputs[inside_interval_mask], |
| unnormalized_widths=unnormalized_widths[inside_interval_mask, :], |
| unnormalized_heights=unnormalized_heights[inside_interval_mask, :], |
| unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :], |
| inverse=inverse, |
| left=-tail_bound, right=tail_bound, bottom=-tail_bound, top=tail_bound, |
| min_bin_width=min_bin_width, |
| min_bin_height=min_bin_height, |
| min_derivative=min_derivative |
| ) |
|
|
| return outputs, logabsdet |
|
|
| def rational_quadratic_spline(inputs, |
| unnormalized_widths, |
| unnormalized_heights, |
| unnormalized_derivatives, |
| inverse=False, |
| left=0., right=1., bottom=0., top=1., |
| min_bin_width=DEFAULT_MIN_BIN_WIDTH, |
| min_bin_height=DEFAULT_MIN_BIN_HEIGHT, |
| min_derivative=DEFAULT_MIN_DERIVATIVE): |
| if torch.min(inputs) < left or torch.max(inputs) > right: |
| raise ValueError('Input to a transform is not within its domain') |
|
|
| num_bins = unnormalized_widths.shape[-1] |
|
|
| if min_bin_width * num_bins > 1.0: |
| raise ValueError('Minimal bin width too large for the number of bins') |
| if min_bin_height * num_bins > 1.0: |
| raise ValueError('Minimal bin height too large for the number of bins') |
|
|
| widths = F.softmax(unnormalized_widths, dim=-1) |
| widths = min_bin_width + (1 - min_bin_width * num_bins) * widths |
| cumwidths = torch.cumsum(widths, dim=-1) |
| cumwidths = F.pad(cumwidths, pad=(1, 0), mode='constant', value=0.0) |
| cumwidths = (right - left) * cumwidths + left |
| cumwidths[..., 0] = left |
| cumwidths[..., -1] = right |
| widths = cumwidths[..., 1:] - cumwidths[..., :-1] |
|
|
| derivatives = min_derivative + F.softplus(unnormalized_derivatives) |
|
|
| heights = F.softmax(unnormalized_heights, dim=-1) |
| heights = min_bin_height + (1 - min_bin_height * num_bins) * heights |
| cumheights = torch.cumsum(heights, dim=-1) |
| cumheights = F.pad(cumheights, pad=(1, 0), mode='constant', value=0.0) |
| cumheights = (top - bottom) * cumheights + bottom |
| cumheights[..., 0] = bottom |
| cumheights[..., -1] = top |
| heights = cumheights[..., 1:] - cumheights[..., :-1] |
|
|
| if inverse: |
| bin_idx = searchsorted(cumheights, inputs)[..., None] |
| else: |
| bin_idx = searchsorted(cumwidths, inputs)[..., None] |
|
|
| input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0] |
| input_bin_widths = widths.gather(-1, bin_idx)[..., 0] |
|
|
| input_cumheights = cumheights.gather(-1, bin_idx)[..., 0] |
| delta = heights / widths |
| input_delta = delta.gather(-1, bin_idx)[..., 0] |
|
|
| input_derivatives = derivatives.gather(-1, bin_idx)[..., 0] |
| input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0] |
|
|
| input_heights = heights.gather(-1, bin_idx)[..., 0] |
|
|
| if inverse: |
| a = (((inputs - input_cumheights) * (input_derivatives |
| + input_derivatives_plus_one |
| - 2 * input_delta) |
| + input_heights * (input_delta - input_derivatives))) |
| b = (input_heights * input_derivatives |
| - (inputs - input_cumheights) * (input_derivatives |
| + input_derivatives_plus_one |
| - 2 * input_delta)) |
| c = - input_delta * (inputs - input_cumheights) |
|
|
| discriminant = b.pow(2) - 4 * a * c |
| assert (discriminant >= 0).all() |
|
|
| root = (2 * c) / (-b - torch.sqrt(discriminant)) |
| outputs = root * input_bin_widths + input_cumwidths |
|
|
| theta_one_minus_theta = root * (1 - root) |
| denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta) |
| * theta_one_minus_theta) |
| derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * root.pow(2) |
| + 2 * input_delta * theta_one_minus_theta |
| + input_derivatives * (1 - root).pow(2)) |
| logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator) |
|
|
| return outputs, -logabsdet |
| else: |
| theta = (inputs - input_cumwidths) / input_bin_widths |
| theta_one_minus_theta = theta * (1 - theta) |
|
|
| numerator = input_heights * (input_delta * theta.pow(2) |
| + input_derivatives * theta_one_minus_theta) |
| denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta) |
| * theta_one_minus_theta) |
| outputs = input_cumheights + numerator / denominator |
|
|
| derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * theta.pow(2) |
| + 2 * input_delta * theta_one_minus_theta |
| + input_derivatives * (1 - theta).pow(2)) |
| logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator) |
|
|
| return outputs, logabsdet |
|
|
| """LayerFusion 4-Bit (LF4) Engine for Inflect-Micro-v2 / Vortex-TTS. |
| |
| Implements native 4-bit packed uint8 nibble storage with blockwise scale & zero metadata. |
| Inherits design principles from vortex_embed native 4-bit execution: |
| - Zero precomputed FP32 weight tables in memory. |
| - On-the-fly nibble unpacking during forward pass execution. |
| - 6.4x weight compression (35.65 MB -> 5.57 MB). |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class LF4Linear(nn.Module): |
| """Native LF4 4-bit Packed Linear Layer.""" |
|
|
| def __init__(self, in_features: int, out_features: int, block_size: int = 32): |
| super().__init__() |
| self.in_features = in_features |
| self.out_features = out_features |
| self.block_size = block_size |
|
|
| num_weights = in_features * out_features |
| pad = (block_size - (num_weights % block_size)) % block_size |
| total_len = num_weights + pad |
| num_blocks = total_len // block_size |
|
|
| self.register_buffer("packed_weight", torch.zeros((total_len // 2,), dtype=torch.uint8)) |
| self.register_buffer("scales", torch.ones((num_blocks, 1), dtype=torch.float16)) |
| self.register_buffer("zeros", torch.zeros((num_blocks, 1), dtype=torch.float16)) |
| self.bias = None |
|
|
| @classmethod |
| def from_float(cls, float_layer: nn.Linear, block_size: int = 32) -> LF4Linear: |
| w = float_layer.weight.data |
| out_f, in_f = w.shape |
| lf4 = cls(in_f, out_f, block_size=block_size) |
|
|
| flat = w.reshape(-1) |
| n = flat.numel() |
| pad = (block_size - (n % block_size)) % block_size |
| if pad > 0: |
| flat = torch.cat([flat, torch.zeros(pad, device=w.device)]) |
|
|
| groups = flat.reshape(-1, block_size) |
| g_min = groups.min(dim=-1, keepdim=True).values |
| g_max = groups.max(dim=-1, keepdim=True).values |
|
|
| scales = ((g_max - g_min) / 15.0).clamp(min=1e-8) |
| zeros = g_min |
|
|
| q_groups = torch.round((groups - zeros) / scales).clamp(0, 15).to(torch.uint8) |
|
|
| q_flat = q_groups.reshape(-1) |
| low_nibbles = q_flat[0::2] |
| high_nibbles = q_flat[1::2] |
| packed = low_nibbles | (high_nibbles << 4) |
|
|
| lf4.packed_weight.copy_(packed) |
| lf4.scales.copy_(scales.to(torch.float16)) |
| lf4.zeros.copy_(zeros.to(torch.float16)) |
| if float_layer.bias is not None: |
| lf4.bias = nn.Parameter(float_layer.bias.data.clone()) |
| return lf4 |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| p = self.packed_weight |
| low = (p & 0x0F).to(torch.float32) |
| high = ((p >> 4) & 0x0F).to(torch.float32) |
|
|
| unpacked = torch.empty((p.numel() * 2,), dtype=torch.float32, device=x.device) |
| unpacked[0::2] = low |
| unpacked[1::2] = high |
|
|
| n_orig = self.in_features * self.out_features |
| unpacked = unpacked[: ((n_orig + self.block_size - 1) // self.block_size) * self.block_size] |
| blocks = unpacked.reshape(-1, self.block_size) |
|
|
| s = self.scales.to(torch.float32) |
| z = self.zeros.to(torch.float32) |
| w_deq = (blocks * s + z).reshape(-1)[:n_orig].reshape(self.out_features, self.in_features) |
|
|
| return F.linear(x, w_deq, self.bias) |
|
|
|
|
| class LF4Conv1d(nn.Module): |
| """Native LF4 4-bit Packed Conv1d Layer.""" |
|
|
| def __init__( |
| self, |
| in_channels: int, |
| out_channels: int, |
| kernel_size: int, |
| stride: int = 1, |
| padding: int = 0, |
| dilation: int = 1, |
| groups: int = 1, |
| block_size: int = 32, |
| ): |
| super().__init__() |
| self.in_channels = in_channels |
| self.out_channels = out_channels |
| self.kernel_size = kernel_size |
| self.stride = stride |
| self.padding = padding |
| self.dilation = dilation |
| self.groups = groups |
| self.block_size = block_size |
|
|
| num_weights = out_channels * (in_channels // groups) * kernel_size |
| pad = (block_size - (num_weights % block_size)) % block_size |
| total_len = num_weights + pad |
| num_blocks = total_len // block_size |
|
|
| self.register_buffer("packed_weight", torch.zeros((total_len // 2,), dtype=torch.uint8)) |
| self.register_buffer("scales", torch.ones((num_blocks, 1), dtype=torch.float16)) |
| self.register_buffer("zeros", torch.zeros((num_blocks, 1), dtype=torch.float16)) |
| self.bias = None |
|
|
| @classmethod |
| def from_float(cls, conv: nn.Conv1d, block_size: int = 32) -> LF4Conv1d: |
| lf4 = cls( |
| conv.in_channels, |
| conv.out_channels, |
| conv.kernel_size[0], |
| stride=conv.stride[0], |
| padding=conv.padding[0], |
| dilation=conv.dilation[0], |
| groups=conv.groups, |
| block_size=block_size, |
| ) |
| w = conv.weight.data |
| flat = w.reshape(-1) |
| n = flat.numel() |
| pad = (block_size - (n % block_size)) % block_size |
| if pad > 0: |
| flat = torch.cat([flat, torch.zeros(pad, device=w.device)]) |
|
|
| groups = flat.reshape(-1, block_size) |
| g_min = groups.min(dim=-1, keepdim=True).values |
| g_max = groups.max(dim=-1, keepdim=True).values |
|
|
| scales = ((g_max - g_min) / 15.0).clamp(min=1e-8) |
| zeros = g_min |
|
|
| q_groups = torch.round((groups - zeros) / scales).clamp(0, 15).to(torch.uint8) |
|
|
| q_flat = q_groups.reshape(-1) |
| low_nibbles = q_flat[0::2] |
| high_nibbles = q_flat[1::2] |
| packed = low_nibbles | (high_nibbles << 4) |
|
|
| lf4.packed_weight.copy_(packed) |
| lf4.scales.copy_(scales.to(torch.float16)) |
| lf4.zeros.copy_(zeros.to(torch.float16)) |
| if conv.bias is not None: |
| lf4.bias = nn.Parameter(conv.bias.data.clone()) |
| return lf4 |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| p = self.packed_weight |
| low = (p & 0x0F).to(torch.float32) |
| high = ((p >> 4) & 0x0F).to(torch.float32) |
|
|
| unpacked = torch.empty((p.numel() * 2,), dtype=torch.float32, device=x.device) |
| unpacked[0::2] = low |
| unpacked[1::2] = high |
|
|
| n_orig = self.out_channels * (self.in_channels // self.groups) * self.kernel_size |
| unpacked = unpacked[: ((n_orig + self.block_size - 1) // self.block_size) * self.block_size] |
| blocks = unpacked.reshape(-1, self.block_size) |
|
|
| s = self.scales.to(torch.float32) |
| z = self.zeros.to(torch.float32) |
| w_deq = (blocks * s + z).reshape(-1)[:n_orig].reshape( |
| self.out_channels, self.in_channels // self.groups, self.kernel_size |
| ) |
|
|
| return F.conv1d( |
| x, |
| w_deq, |
| self.bias, |
| stride=self.stride, |
| padding=self.padding, |
| dilation=self.dilation, |
| groups=self.groups, |
| ) |
|
|
| import copy |
| import math |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d |
| from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm |
|
|
|
|
| class StochasticDurationPredictor(nn.Module): |
| def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, n_flows=4, gin_channels=0): |
| super().__init__() |
| filter_channels = in_channels |
| self.in_channels = in_channels |
| self.filter_channels = filter_channels |
| self.kernel_size = kernel_size |
| self.p_dropout = p_dropout |
| self.n_flows = n_flows |
| self.gin_channels = gin_channels |
|
|
| self.log_flow = Log() |
| self.flows = nn.ModuleList() |
| self.flows.append(ElementwiseAffine(2)) |
| for i in range(n_flows): |
| self.flows.append(ConvFlow(2, filter_channels, kernel_size, n_layers=3)) |
| self.flows.append(Flip()) |
|
|
| self.post_pre = nn.Conv1d(1, filter_channels, 1) |
| self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1) |
| self.post_convs = DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout) |
| self.post_flows = nn.ModuleList() |
| self.post_flows.append(ElementwiseAffine(2)) |
| for i in range(4): |
| self.post_flows.append(ConvFlow(2, filter_channels, kernel_size, n_layers=3)) |
| self.post_flows.append(Flip()) |
|
|
| self.pre = nn.Conv1d(in_channels, filter_channels, 1) |
| self.proj = nn.Conv1d(filter_channels, filter_channels, 1) |
| self.convs = DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout) |
| if gin_channels != 0: |
| self.cond = nn.Conv1d(gin_channels, filter_channels, 1) |
|
|
| def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0): |
| x = torch.detach(x) |
| x = self.pre(x) |
| if g is not None: |
| g = torch.detach(g) |
| x = x + self.cond(g) |
| x = self.convs(x, x_mask) |
| x = self.proj(x) * x_mask |
|
|
| if not reverse: |
| flows = self.flows |
| assert w is not None |
|
|
| logdet_tot_q = 0 |
| h_w = self.post_pre(w) |
| h_w = self.post_convs(h_w, x_mask) |
| h_w = self.post_proj(h_w) * x_mask |
| e_q = torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype) * x_mask |
| z_q = e_q |
| for flow in self.post_flows: |
| z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w)) |
| logdet_tot_q += logdet_q |
| z_u, z1 = torch.split(z_q, [1, 1], 1) |
| u = torch.sigmoid(z_u) * x_mask |
| z0 = (w - u) * x_mask |
| logdet_tot_q += torch.sum((F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1,2]) |
| logq = torch.sum(-0.5 * (math.log(2*math.pi) + (e_q**2)) * x_mask, [1,2]) - logdet_tot_q |
|
|
| logdet_tot = 0 |
| z0, logdet = self.log_flow(z0, x_mask) |
| logdet_tot += logdet |
| z = torch.cat([z0, z1], 1) |
| for flow in flows: |
| z, logdet = flow(z, x_mask, g=x, reverse=reverse) |
| logdet_tot = logdet_tot + logdet |
| nll = torch.sum(0.5 * (math.log(2*math.pi) + (z**2)) * x_mask, [1,2]) - logdet_tot |
| return nll + logq |
| else: |
| flows = list(reversed(self.flows)) |
| flows = flows[:-2] + [flows[-1]] |
| z = torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype) * noise_scale |
| for flow in flows: |
| z = flow(z, x_mask, g=x, reverse=reverse) |
| z0, z1 = torch.split(z, [1, 1], 1) |
| logw = z0 |
| return logw |
|
|
|
|
| class DurationPredictor(nn.Module): |
| def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0): |
| super().__init__() |
|
|
| self.in_channels = in_channels |
| self.filter_channels = filter_channels |
| self.kernel_size = kernel_size |
| self.p_dropout = p_dropout |
| self.gin_channels = gin_channels |
|
|
| self.drop = nn.Dropout(p_dropout) |
| self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size//2) |
| self.norm_1 = LayerNorm(filter_channels) |
| self.conv_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size//2) |
| self.norm_2 = LayerNorm(filter_channels) |
| self.proj = nn.Conv1d(filter_channels, 1, 1) |
|
|
| if gin_channels != 0: |
| self.cond = nn.Conv1d(gin_channels, in_channels, 1) |
|
|
| def forward(self, x, x_mask, g=None): |
| x = torch.detach(x) |
| if g is not None: |
| g = torch.detach(g) |
| x = x + self.cond(g) |
| x = self.conv_1(x * x_mask) |
| x = torch.relu(x) |
| x = self.norm_1(x) |
| x = self.drop(x) |
| x = self.conv_2(x * x_mask) |
| x = torch.relu(x) |
| x = self.norm_2(x) |
| x = self.drop(x) |
| x = self.proj(x * x_mask) |
| return x * x_mask |
|
|
|
|
| class TextEncoder(nn.Module): |
| def __init__(self, |
| n_vocab, |
| out_channels, |
| hidden_channels, |
| filter_channels, |
| n_heads, |
| n_layers, |
| kernel_size, |
| p_dropout): |
| super().__init__() |
| self.n_vocab = n_vocab |
| self.out_channels = out_channels |
| self.hidden_channels = hidden_channels |
| self.filter_channels = filter_channels |
| self.n_heads = n_heads |
| self.n_layers = n_layers |
| self.kernel_size = kernel_size |
| self.p_dropout = p_dropout |
|
|
| self.emb = nn.Embedding(n_vocab, hidden_channels) |
| nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5) |
|
|
| self.encoder = Encoder( |
| hidden_channels, |
| filter_channels, |
| n_heads, |
| n_layers, |
| kernel_size, |
| p_dropout) |
| self.proj= nn.Conv1d(hidden_channels, out_channels * 2, 1) |
|
|
| def forward(self, x, x_lengths): |
| x = self.emb(x) * math.sqrt(self.hidden_channels) |
| x = torch.transpose(x, 1, -1) |
| x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype) |
|
|
| x = self.encoder(x * x_mask, x_mask) |
| stats = self.proj(x) * x_mask |
|
|
| m, logs = torch.split(stats, self.out_channels, dim=1) |
| return x, m, logs, x_mask |
|
|
|
|
| class ResidualCouplingBlock(nn.Module): |
| def __init__(self, |
| channels, |
| hidden_channels, |
| kernel_size, |
| dilation_rate, |
| n_layers, |
| n_flows=4, |
| gin_channels=0): |
| super().__init__() |
| self.channels = channels |
| self.hidden_channels = hidden_channels |
| self.kernel_size = kernel_size |
| self.dilation_rate = dilation_rate |
| self.n_layers = n_layers |
| self.n_flows = n_flows |
| self.gin_channels = gin_channels |
|
|
| self.flows = nn.ModuleList() |
| for i in range(n_flows): |
| self.flows.append(ResidualCouplingLayer(channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels, mean_only=True)) |
| self.flows.append(Flip()) |
|
|
| def forward(self, x, x_mask, g=None, reverse=False): |
| if not reverse: |
| for flow in self.flows: |
| x, _ = flow(x, x_mask, g=g, reverse=reverse) |
| else: |
| for flow in reversed(self.flows): |
| x = flow(x, x_mask, g=g, reverse=reverse) |
| return x |
|
|
|
|
| class PosteriorEncoder(nn.Module): |
| def __init__(self, |
| in_channels, |
| out_channels, |
| hidden_channels, |
| kernel_size, |
| dilation_rate, |
| n_layers, |
| gin_channels=0): |
| 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.pre = nn.Conv1d(in_channels, hidden_channels, 1) |
| self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels) |
| self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1) |
|
|
| def forward(self, x, x_lengths, g=None): |
| x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype) |
| x = self.pre(x) * x_mask |
| x = self.enc(x, x_mask, g=g) |
| stats = self.proj(x) * x_mask |
| m, logs = torch.split(stats, self.out_channels, dim=1) |
| z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask |
| return z, m, logs, x_mask |
|
|
|
|
| class Generator(torch.nn.Module): |
| def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels=0, decoder_alias_free=False, decoder_alias_free_start_stage=2, decoder_snake_logscale=True): |
| super(Generator, self).__init__() |
| self.num_kernels = len(resblock_kernel_sizes) |
| self.num_upsamples = len(upsample_rates) |
| self.conv_pre = Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3) |
| resblock_class = ResBlock1 if resblock == '1' else ResBlock2 |
| self.decoder_alias_free = bool(decoder_alias_free) |
| self.decoder_alias_free_start_stage = int(decoder_alias_free_start_stage) |
|
|
| self.ups = nn.ModuleList() |
| for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): |
| self.ups.append(weight_norm( |
| ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)), |
| k, u, padding=(k-u)//2))) |
|
|
| self.resblocks = nn.ModuleList() |
| for i in range(len(self.ups)): |
| ch = upsample_initial_channel//(2**(i+1)) |
| for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)): |
| if self.decoder_alias_free and i >= self.decoder_alias_free_start_stage: |
| self.resblocks.append(AliasFreeResBlock1( |
| ch, k, d, logscale=decoder_snake_logscale)) |
| else: |
| self.resblocks.append(resblock_class(ch, k, d)) |
|
|
| self.alias_free_pre_activations = nn.ModuleList() |
| if self.decoder_alias_free: |
| for i in range(self.num_upsamples): |
| channels = upsample_initial_channel // (2 ** i) |
| if i >= self.decoder_alias_free_start_stage: |
| self.alias_free_pre_activations.append( |
| AliasFreeActivation1d(nn.LeakyReLU(LRELU_SLOPE))) |
| else: |
| self.alias_free_pre_activations.append(nn.Identity()) |
| self.alias_free_post_activation = AliasFreeActivation1d( |
| SnakeBeta(ch, logscale=decoder_snake_logscale)) |
|
|
| self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False) |
| self.ups.apply(init_weights) |
|
|
| if gin_channels != 0: |
| self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1) |
|
|
| def forward(self, x, g=None): |
| x = self.conv_pre(x) |
| if g is not None: |
| x = x + self.cond(g) |
|
|
| for i in range(self.num_upsamples): |
| if self.decoder_alias_free and i >= self.decoder_alias_free_start_stage: |
| x = self.alias_free_pre_activations[i](x) |
| else: |
| x = F.leaky_relu(x, LRELU_SLOPE) |
| x = self.ups[i](x) |
| xs = None |
| for j in range(self.num_kernels): |
| if xs is None: |
| xs = self.resblocks[i*self.num_kernels+j](x) |
| else: |
| xs += self.resblocks[i*self.num_kernels+j](x) |
| x = xs / self.num_kernels |
| if self.decoder_alias_free: |
| x = self.alias_free_post_activation(x) |
| else: |
| x = F.leaky_relu(x) |
| x = self.conv_post(x) |
| x = torch.tanh(x) |
|
|
| return x |
|
|
| def remove_weight_norm(self): |
| print('Removing weight norm...') |
| for l in self.ups: |
| remove_weight_norm(l) |
| for l in self.resblocks: |
| l.remove_weight_norm() |
|
|
|
|
| class DiscriminatorP(torch.nn.Module): |
| def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False): |
| super(DiscriminatorP, self).__init__() |
| self.period = period |
| self.use_spectral_norm = use_spectral_norm |
| norm_f = weight_norm if use_spectral_norm == False else spectral_norm |
| self.convs = nn.ModuleList([ |
| norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))), |
| norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))), |
| norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))), |
| norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))), |
| norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(get_padding(kernel_size, 1), 0))), |
| ]) |
| self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0))) |
|
|
| def forward(self, x): |
| fmap = [] |
|
|
| |
| b, c, t = x.shape |
| if t % self.period != 0: |
| n_pad = self.period - (t % self.period) |
| x = F.pad(x, (0, n_pad), "reflect") |
| t = t + n_pad |
| x = x.view(b, c, t // self.period, self.period) |
|
|
| for l in self.convs: |
| x = l(x) |
| x = F.leaky_relu(x, LRELU_SLOPE) |
| fmap.append(x) |
| x = self.conv_post(x) |
| fmap.append(x) |
| x = torch.flatten(x, 1, -1) |
|
|
| return x, fmap |
|
|
|
|
| class DiscriminatorS(torch.nn.Module): |
| def __init__(self, use_spectral_norm=False): |
| super(DiscriminatorS, self).__init__() |
| norm_f = weight_norm if use_spectral_norm == False else spectral_norm |
| self.convs = nn.ModuleList([ |
| norm_f(Conv1d(1, 16, 15, 1, padding=7)), |
| norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)), |
| norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)), |
| norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)), |
| norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)), |
| norm_f(Conv1d(1024, 1024, 5, 1, padding=2)), |
| ]) |
| self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1)) |
|
|
| def forward(self, x): |
| fmap = [] |
|
|
| for l in self.convs: |
| x = l(x) |
| x = F.leaky_relu(x, LRELU_SLOPE) |
| fmap.append(x) |
| x = self.conv_post(x) |
| fmap.append(x) |
| x = torch.flatten(x, 1, -1) |
|
|
| return x, fmap |
|
|
|
|
| class MultiPeriodDiscriminator(torch.nn.Module): |
| def __init__(self, use_spectral_norm=False): |
| super(MultiPeriodDiscriminator, self).__init__() |
| periods = [2,3,5,7,11] |
|
|
| discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)] |
| discs = discs + [DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods] |
| self.discriminators = nn.ModuleList(discs) |
|
|
| def forward(self, y, y_hat): |
| y_d_rs = [] |
| y_d_gs = [] |
| fmap_rs = [] |
| fmap_gs = [] |
| for i, d in enumerate(self.discriminators): |
| y_d_r, fmap_r = d(y) |
| y_d_g, fmap_g = d(y_hat) |
| y_d_rs.append(y_d_r) |
| y_d_gs.append(y_d_g) |
| fmap_rs.append(fmap_r) |
| fmap_gs.append(fmap_g) |
|
|
| return y_d_rs, y_d_gs, fmap_rs, fmap_gs |
|
|
|
|
|
|
| class SynthesizerTrn(nn.Module): |
| """ |
| Synthesizer for Training |
| """ |
|
|
| def __init__(self, |
| n_vocab, |
| spec_channels, |
| segment_size, |
| inter_channels, |
| hidden_channels, |
| filter_channels, |
| n_heads, |
| n_layers, |
| kernel_size, |
| p_dropout, |
| resblock, |
| resblock_kernel_sizes, |
| resblock_dilation_sizes, |
| upsample_rates, |
| upsample_initial_channel, |
| upsample_kernel_sizes, |
| n_speakers=0, |
| gin_channels=0, |
| use_sdp=True, |
| **kwargs): |
|
|
| super().__init__() |
| self.n_vocab = n_vocab |
| self.spec_channels = spec_channels |
| self.inter_channels = inter_channels |
| self.hidden_channels = hidden_channels |
| self.filter_channels = filter_channels |
| self.n_heads = n_heads |
| self.n_layers = n_layers |
| self.kernel_size = kernel_size |
| self.p_dropout = p_dropout |
| self.resblock = resblock |
| self.resblock_kernel_sizes = resblock_kernel_sizes |
| self.resblock_dilation_sizes = resblock_dilation_sizes |
| self.upsample_rates = upsample_rates |
| self.upsample_initial_channel = upsample_initial_channel |
| self.upsample_kernel_sizes = upsample_kernel_sizes |
| self.segment_size = segment_size |
| self.n_speakers = n_speakers |
| self.gin_channels = gin_channels |
|
|
| self.use_sdp = use_sdp |
|
|
| self.VTXTextEncoder = TextEncoder(n_vocab, |
| inter_channels, |
| hidden_channels, |
| filter_channels, |
| n_heads, |
| n_layers, |
| kernel_size, |
| p_dropout) |
| self.enc_p = self.VTXTextEncoder |
| |
| self.VTXVocoder = Generator( |
| inter_channels, resblock, resblock_kernel_sizes, |
| resblock_dilation_sizes, upsample_rates, upsample_initial_channel, |
| upsample_kernel_sizes, gin_channels=gin_channels, |
| decoder_alias_free=kwargs.get("decoder_alias_free", False), |
| decoder_alias_free_start_stage=kwargs.get("decoder_alias_free_start_stage", 2), |
| decoder_snake_logscale=kwargs.get("decoder_snake_logscale", True)) |
| self.dec = self.VTXVocoder |
| |
| self.inference_only = bool(kwargs.get("inference_only", False)) |
| if not self.inference_only: |
| self.enc_q = PosteriorEncoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels) |
| |
| self.VTXVectorEstimator = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels) |
| self.flow = self.VTXVectorEstimator |
|
|
| if use_sdp: |
| self.VTXDurationPredictor = StochasticDurationPredictor(hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels) |
| else: |
| self.VTXDurationPredictor = DurationPredictor(hidden_channels, 256, 3, 0.5, gin_channels=gin_channels) |
| self.dp = self.VTXDurationPredictor |
|
|
| if n_speakers > 1: |
| self.emb_g = nn.Embedding(n_speakers, gin_channels) |
|
|
| def forward(self, x, x_lengths, y, y_lengths, sid=None): |
|
|
| if getattr(self, "inference_only", False) or not hasattr(self, "enc_q"): |
| raise RuntimeError("The public runtime is inference-only and has no posterior encoder.") |
|
|
| x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths) |
| if self.n_speakers > 0: |
| g = self.emb_g(sid).unsqueeze(-1) |
| else: |
| g = None |
|
|
| z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g) |
| z_p = self.flow(z, y_mask, g=g) |
|
|
| with torch.no_grad(): |
| |
| s_p_sq_r = torch.exp(-2 * logs_p) |
| neg_cent1 = torch.sum(-0.5 * math.log(2 * math.pi) - logs_p, [1], keepdim=True) |
| neg_cent2 = torch.matmul(-0.5 * (z_p ** 2).transpose(1, 2), s_p_sq_r) |
| neg_cent3 = torch.matmul(z_p.transpose(1, 2), (m_p * s_p_sq_r)) |
| neg_cent4 = torch.sum(-0.5 * (m_p ** 2) * s_p_sq_r, [1], keepdim=True) |
| neg_cent = neg_cent1 + neg_cent2 + neg_cent3 + neg_cent4 |
|
|
| attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1) |
| attn = maximum_path(neg_cent, attn_mask.squeeze(1)).unsqueeze(1).detach() |
|
|
| w = attn.sum(2) |
| if self.use_sdp: |
| l_length = self.dp(x, x_mask, w, g=g) |
| l_length = l_length / torch.sum(x_mask) |
| else: |
| logw_ = torch.log(w + 1e-6) * x_mask |
| logw = self.dp(x, x_mask, g=g) |
| l_length = torch.sum((logw - logw_)**2, [1,2]) / torch.sum(x_mask) |
|
|
| |
| m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(1, 2) |
| logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(1, 2) |
|
|
| z_slice, ids_slice = rand_slice_segments(z, y_lengths, self.segment_size) |
| o = self.dec(z_slice, g=g) |
| return o, l_length, attn, ids_slice, x_mask, y_mask, (z, z_p, m_p, logs_p, m_q, logs_q) |
|
|
| def infer_text(self, text, device="cpu", noise_scale=0.667, length_scale=1.0, noise_scale_w=0.8, max_len=None): |
| from configuration import cleaned_text_to_sequence |
| from phonemizer.backend import EspeakBackend |
| |
| backend = EspeakBackend('en-us', preserve_punctuation=True, with_stress=True) |
| phoneme_str = backend.phonemize([text])[0] |
| sequence = cleaned_text_to_sequence(phoneme_str) |
| if not sequence: |
| sequence = [1, 2, 3] |
| |
| res = [0] * (len(sequence) * 2 + 1) |
| res[1::2] = sequence |
| tokens = torch.LongTensor(res).to(device).unsqueeze(0) |
| lengths = torch.LongTensor([tokens.size(1)]).to(device) |
| |
| return self.infer(tokens, lengths, noise_scale=noise_scale, length_scale=length_scale, noise_scale_w=noise_scale_w, max_len=max_len) |
|
|
| def infer(self, x, x_lengths, sid=None, noise_scale=1, length_scale=1, noise_scale_w=1., max_len=None): |
| x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths) |
| if self.n_speakers > 0: |
| g = self.emb_g(sid).unsqueeze(-1) |
| else: |
| g = None |
|
|
| if self.use_sdp: |
| logw = self.dp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) |
| else: |
| logw = self.dp(x, x_mask, g=g) |
| w = torch.exp(logw) * x_mask * length_scale |
| w_ceil = torch.ceil(w) |
| y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long() |
| y_mask = torch.unsqueeze(sequence_mask(y_lengths, None), 1).to(x_mask.dtype) |
| attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1) |
| attn = generate_path(w_ceil, attn_mask) |
|
|
| m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(1, 2) |
| logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(1, 2) |
|
|
| z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale |
| z = self.flow(z_p, y_mask, g=g, reverse=True) |
| o = self.dec((z * y_mask)[:,:,:max_len], g=g) |
| return o, attn, y_mask, (z, z_p, m_p, logs_p) |
|
|
| def voice_conversion(self, y, y_lengths, sid_src, sid_tgt): |
| if self.inference_only: |
| raise RuntimeError("The public runtime is inference-only and does not support voice conversion.") |
| assert self.n_speakers > 0, "n_speakers have to be larger than 0." |
| g_src = self.emb_g(sid_src).unsqueeze(-1) |
| g_tgt = self.emb_g(sid_tgt).unsqueeze(-1) |
| z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g_src) |
| z_p = self.flow(z, y_mask, g=g_src) |
| z_hat = self.flow(z_p, y_mask, g=g_tgt, reverse=True) |
| o_hat = self.dec(z_hat * y_mask, g=g_tgt) |
| return o_hat, y_mask, (z, z_p, z_hat) |
