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): # reshape [b, d, t] -> [b, n_h, t, d_k] 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) # [b, n_h, t_t, t_s] 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) # [b, n_h, t_t, d_k] -> [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 first before slice to avoid using cond ops. 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() # Concat columns of pad to shift from relative to absolute indexing. x = F.pad(x, convert_pad_shape([[0,0],[0,0],[0,0],[0,1]])) # Concat extra elements so to add up to shape (len+1, 2*len-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]])) # Reshape and slice out the padded elements. 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() # padd along column 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)]) # add 0's in the beginning that will skew the elements after reshape 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) # last one is not necessary 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) # [b, cx?, t] -> [b, c, t, ?] 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 # it needs to be removed from future version. 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 # [b] else: flows = list(reversed(self.flows)) flows = flows[:-2] + [flows[-1]] # remove a useless vflow 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) # [b, t, h] x = torch.transpose(x, 1, -1) # [b, h, t] 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 = [] # 1d to 2d b, c, t = x.shape if t % self.period != 0: # pad first 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) # [b, h, 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(): # negative cross-entropy s_p_sq_r = torch.exp(-2 * logs_p) # [b, d, t] neg_cent1 = torch.sum(-0.5 * math.log(2 * math.pi) - logs_p, [1], keepdim=True) # [b, 1, t_s] neg_cent2 = torch.matmul(-0.5 * (z_p ** 2).transpose(1, 2), s_p_sq_r) # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s] neg_cent3 = torch.matmul(z_p.transpose(1, 2), (m_p * s_p_sq_r)) # [b, t_t, d] x [b, d, t_s] = [b, t_t, t_s] neg_cent4 = torch.sum(-0.5 * (m_p ** 2) * s_p_sq_r, [1], keepdim=True) # [b, 1, t_s] 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) # for averaging # expand prior 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) # [b, h, 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) # [b, t', t], [b, t, d] -> [b, d, t'] logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(1, 2) # [b, t', t], [b, t, d] -> [b, d, t'] 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)