| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch import autocast |
|
|
| from config import config |
|
|
| from .modules import GST, AttentionBlock, mySequential, normalization |
|
|
|
|
| def timestep_embedding(timesteps, dim, max_period=10000): |
| """ |
| Create sinusoidal timestep embeddings. |
| |
| :param timesteps: a 1-D Tensor of N indices, one per batch element. |
| These may be fractional. |
| :param dim: the dimension of the output. |
| :param max_period: controls the minimum frequency of the embeddings. |
| :return: an [N x dim] Tensor of positional embeddings. |
| """ |
| half = dim // 2 |
| freqs = torch.exp( |
| -math.log(max_period) |
| * torch.arange(start=0, end=half, dtype=torch.float32) |
| / half |
| ).to(device=timesteps.device) |
| args = timesteps[:, None].float() * freqs[None] |
| embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) |
| if dim % 2: |
| embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) |
| return embedding |
|
|
|
|
| class TimestepBlock(nn.Module): |
| def forward(self, x, emb): |
| """ |
| Apply the module to `x` given `emb` timestep embeddings. |
| """ |
|
|
|
|
| class TimestepEmbedSequential(nn.Sequential, TimestepBlock): |
| def forward(self, x, emb): |
| for layer in self: |
| if isinstance(layer, TimestepBlock): |
| x = layer(x, emb) |
| else: |
| x = layer(x) |
| return x |
|
|
|
|
| class QuartzNetBlock(TimestepBlock): |
| """Similar to Resnet block with Batchnorm and dropout, and using Separable conv in the middle. |
| if its the last layer,set se = False and separable = False, and use a projection layer on top of this. |
| """ |
|
|
| def __init__( |
| self, |
| nin, |
| nout, |
| emb_channels, |
| kernel_size=3, |
| dropout=0.1, |
| R=1, |
| se=True, |
| ratio=8, |
| separable=False, |
| bias=True, |
| use_scale_shift_norm=True, |
| ): |
| super(QuartzNetBlock, self).__init__() |
| self.use_scale_shift_norm = use_scale_shift_norm |
| self.se = se |
| self.in_layers = mySequential( |
| nn.Conv1d(nin, nout, kernel_size=1, padding="same", bias=bias), |
| nn.SiLU(), |
| normalization(nout), |
| ) |
|
|
| if nin == nout: |
| self.residual = nn.Identity() |
| else: |
| self.residual = nn.Conv1d( |
| nin, nout, kernel_size=1, padding="same", bias=bias |
| ) |
|
|
| nin = nout |
| self.model = nn.Sequential( |
| nn.Conv1d(nin, nout, kernel_size, padding="same"), |
| nn.SiLU(), |
| normalization(nout), |
| nn.Dropout(p=dropout), |
| ) |
|
|
| self.emb_layers = nn.Sequential( |
| nn.Linear( |
| emb_channels, |
| 2 * nout if use_scale_shift_norm else nout, |
| ), |
| nn.SiLU(), |
| ) |
|
|
| def forward(self, x, emb, mask=None): |
| x_new = self.in_layers(x) |
| emb = self.emb_layers(emb) |
| while len(emb.shape) < len(x_new.shape): |
| emb = emb[..., None] |
| scale, shift = torch.chunk(emb, 2, dim=1) |
| x_new = x_new * (1 + scale) + shift |
| y = self.model(x_new) |
|
|
| return y + self.residual(x) |
|
|
|
|
| class QuartzAttn(TimestepBlock): |
| def __init__(self, model_channels, dropout, num_heads): |
| super().__init__() |
| self.resblk = QuartzNetBlock( |
| model_channels, |
| model_channels, |
| model_channels, |
| dropout=dropout, |
| use_scale_shift_norm=True, |
| ) |
| self.attn = AttentionBlock( |
| model_channels, num_heads, relative_pos_embeddings=True |
| ) |
|
|
| def forward(self, x, time_emb): |
| y = self.resblk(x, time_emb) |
| return self.attn(y) |
|
|
|
|
| class QuartzNet9x5(nn.Module): |
| def __init__(self, model_channels, num_heads, dropout=0.1, enable_fp16=False): |
| super(QuartzNet9x5, self).__init__() |
| self.enable_fp16 = enable_fp16 |
| kernels = [3] * 10 |
| quartznet = [] |
| attn = [] |
| for i in kernels: |
| quartznet.append( |
| QuartzNetBlock( |
| model_channels, |
| model_channels, |
| model_channels, |
| kernel_size=i, |
| dropout=dropout, |
| R=5, |
| se=True, |
| ) |
| ) |
| attn.append( |
| AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True) |
| ) |
|
|
| self.quartznet = nn.ModuleList(quartznet) |
| self.attn = nn.ModuleList(attn) |
| self.conv2 = nn.ModuleList( |
| [ |
| QuartzNetBlock( |
| model_channels, |
| model_channels, |
| model_channels, |
| kernel_size=3, |
| dropout=dropout, |
| R=3, |
| separable=False, |
| ) |
| for i in range(3) |
| ] |
| ) |
| self.conv3 = nn.Sequential( |
| nn.Conv1d(model_channels, model_channels, 3, padding="same"), |
| nn.SiLU(), |
| normalization(model_channels), |
| nn.Conv1d(model_channels, 100, 1, padding="same"), |
| ) |
|
|
| def forward(self, x, time_emb): |
| for n, (layer, attn) in enumerate(zip(self.quartznet, self.attn)): |
| x = layer(x, time_emb) |
| x = attn(x) |
| for layer in self.conv2: |
| x = layer(x, time_emb) |
|
|
| x = self.conv3(x) |
| return x |
|
|
|
|
| class DiffModel(nn.Module): |
| def __init__( |
| self, |
| input_channels=80, |
| output_channels=160, |
| model_channels=256, |
| num_heads=8, |
| dropout=0.1, |
| num_layers=8, |
| multispeaker=True, |
| style_tokens=100, |
| enable_fp16=False, |
| condition_free_per=0.1, |
| training=False, |
| ar_active=False, |
| in_latent_channels=10004, |
| ): |
| super().__init__() |
| self.input_channels = input_channels |
| self.model_channels = model_channels |
| self.output_channels = output_channels |
| self.num_heads = num_heads |
| self.dropout = dropout |
| self.num_layers = num_layers |
| self.enable_fp16 = enable_fp16 |
| self.condition_free_per = condition_free_per |
| self.training = training |
| self.multispeaker = multispeaker |
| self.ar_active = ar_active |
| self.in_latent_channels = in_latent_channels |
|
|
| if not self.ar_active: |
| self.code_emb = nn.Embedding( |
| config.semantic_model_centroids + 1, model_channels |
| ) |
| self.code_converter = mySequential( |
| AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), |
| AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), |
| AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), |
| ) |
| else: |
| self.code_converter = mySequential( |
| nn.Conv1d( |
| self.in_latent_channels, model_channels, 3, padding=1, bias=True |
| ), |
| AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), |
| AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), |
| AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), |
| AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), |
| ) |
| if self.multispeaker: |
| self.GST = GST( |
| model_channels, style_tokens, num_heads, in_channels=input_channels |
| ) |
|
|
| self.code_norm = normalization(model_channels) |
| self.time_norm = normalization(model_channels) |
| self.code_time_norm = normalization(model_channels) |
|
|
| self.time_embed = mySequential( |
| nn.Linear(model_channels, model_channels), |
| nn.SiLU(), |
| nn.Linear(model_channels, model_channels), |
| ) |
|
|
| self.input_block = nn.Conv1d(input_channels, model_channels, 3, 1, 1, bias=True) |
| self.unconditioned_embedding = nn.Parameter(torch.randn(1, model_channels, 1)) |
| self.integrating_conv = nn.Conv1d( |
| model_channels * 2, model_channels, kernel_size=1 |
| ) |
|
|
| self.code_time = TimestepEmbedSequential( |
| QuartzAttn(model_channels, dropout, num_heads), |
| QuartzAttn(model_channels, dropout, num_heads), |
| QuartzAttn(model_channels, dropout, num_heads), |
| ) |
|
|
| self.layers = QuartzNet9x5( |
| model_channels, num_heads, self.enable_fp16, self.dropout |
| ) |
|
|
| def get_speaker_latent(self, ref_mels): |
| ref_mels = ref_mels.unsqueeze(1) if len(ref_mels.shape) == 3 else ref_mels |
|
|
| conds = [] |
| for j in range(ref_mels.shape[1]): |
| conds.append(self.GST(ref_mels[:, j, :, :])) |
|
|
| conds = torch.cat(conds, dim=-1) |
| conds = conds.mean(dim=-1) |
|
|
| return conds.unsqueeze(2) |
|
|
| def forward( |
| self, |
| x, |
| t, |
| code_emb, |
| ref_clips=None, |
| speaker_latents=None, |
| conditioning_free=False, |
| ): |
| time_embed = self.time_norm( |
| self.time_embed( |
| timestep_embedding(t.unsqueeze(-1), self.model_channels) |
| ).permute(0, 2, 1) |
| ).squeeze(2) |
| if conditioning_free: |
| code_embed = self.unconditioned_embedding.repeat(x.shape[0], 1, x.shape[-1]) |
| else: |
| if not self.ar_active: |
| code_embed = self.code_norm( |
| self.code_converter(self.code_emb(code_emb).permute(0, 2, 1)) |
| ) |
| else: |
| code_embed = self.code_norm(self.code_converter(code_emb)) |
| if self.multispeaker: |
| assert speaker_latents is not None or ref_clips is not None |
| if ref_clips is not None: |
| speaker_latents = self.get_speaker_latent(ref_clips) |
| cond_scale, cond_shift = torch.chunk(speaker_latents, 2, dim=1) |
| code_embed = code_embed * (1 + cond_scale) + cond_shift |
|
|
| if self.training and self.condition_free_per > 0: |
| unconditioned_batches = ( |
| torch.rand((code_embed.shape[0], 1, 1), device=code_embed.device) |
| < self.condition_free_per |
| ) |
| code_embed = torch.where( |
| unconditioned_batches, |
| self.unconditioned_embedding.repeat(code_embed.shape[0], 1, 1), |
| code_embed, |
| ) |
|
|
| expanded_code_emb = F.interpolate(code_embed, size=x.shape[-1], mode="linear") |
|
|
| x_cond = self.code_time_norm(self.code_time(expanded_code_emb, time_embed)) |
|
|
| x = self.input_block(x) |
| x = torch.cat([x, x_cond], dim=1) |
| x = self.integrating_conv(x) |
| out = self.layers(x, time_embed) |
|
|
| return out |
|
|
