|
import torch |
|
import torch.distributed |
|
import numpy as np |
|
import logging |
|
import math |
|
import copy |
|
import numpy as np |
|
import scipy |
|
import torch |
|
import librosa |
|
|
|
from typing import Optional, Tuple |
|
from torch import nn, view_as_real, view_as_complex |
|
from torch import nn |
|
from torch.nn import functional as F |
|
from torch.nn.utils import weight_norm, remove_weight_norm |
|
from torchaudio.functional.functional import _hz_to_mel, _mel_to_hz |
|
from transformers.activations import ACT2FN |
|
from dataclasses import dataclass |
|
from transformers.modeling_outputs import ModelOutput |
|
from transformers import WhisperModel |
|
|
|
|
|
|
|
def sinusoids(length, channels, max_timescale=10000): |
|
"""Returns sinusoidal waves for positional embedding""" |
|
assert channels % 2 == 0 |
|
log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1) |
|
inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2)) |
|
scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :] |
|
return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1) |
|
|
|
|
|
def get_sequence_mask(inputs, inputs_length): |
|
if inputs.dim() == 3: |
|
bsz, tgt_len, _ = inputs.size() |
|
else: |
|
bsz, tgt_len = inputs_length.shape[0], torch.max(inputs_length) |
|
sequence_mask = torch.arange(0, tgt_len).to(inputs.device) |
|
sequence_mask = torch.lt(sequence_mask, inputs_length.reshape(bsz, 1)).view(bsz, tgt_len, 1) |
|
return sequence_mask |
|
|
|
|
|
class RMSNorm(nn.Module): |
|
def __init__(self, hidden_size, eps=1e-6): |
|
super().__init__() |
|
self.weight = nn.Parameter(torch.ones(hidden_size)) |
|
self.variance_epsilon = eps |
|
|
|
def forward(self, hidden_states): |
|
variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) |
|
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) |
|
if self.weight.dtype in [torch.float16, torch.bfloat16]: |
|
hidden_states = hidden_states.to(self.weight.dtype) |
|
return self.weight * hidden_states |
|
|
|
|
|
class VarLenAttention(nn.Module): |
|
def __init__(self, embed_dim, num_heads, causal=False, dropout=0.0): |
|
""" |
|
Initialize variable-length attention module. |
|
|
|
Parameters: |
|
embed_dim (int): Embedding dimension (model's hidden dimension) |
|
num_heads (int): Number of attention heads |
|
causal (bool): Whether to enable causal attention (only attend to current and previous positions) |
|
dropout (float): Attention dropout probability |
|
""" |
|
super().__init__() |
|
self.embed_dim = embed_dim |
|
self.num_heads = num_heads |
|
self.head_dim = embed_dim // num_heads |
|
assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads" |
|
self.causal = causal |
|
self.dropout = nn.Dropout(dropout) |
|
self.scaling = self.head_dim ** -0.5 |
|
|
|
|
|
self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False) |
|
self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True) |
|
self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True) |
|
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True) |
|
|
|
def _create_attention_mask(self, seq_len, max_len, device, dtype): |
|
""" |
|
Create attention mask supporting variable-length sequences and causality. |
|
|
|
Parameters: |
|
seq_len (torch.Tensor): Sequence length for each sample, shape [bsz] |
|
max_len (int): Maximum sequence length in the batch |
|
device: Device for tensor creation |
|
dtype: Data type for mask values |
|
|
|
Returns: |
|
mask (torch.Tensor): Attention mask, shape [bsz, 1, max_len, max_len], invalid positions set to minimum value |
|
""" |
|
bsz = seq_len.size(0) |
|
|
|
mask = torch.ones(bsz, 1, max_len, max_len, device=device, dtype=dtype) |
|
|
|
|
|
seq_indices = torch.arange(max_len, device=device).unsqueeze(0) |
|
seq_len_expanded = seq_len.unsqueeze(1) |
|
|
|
|
|
valid_mask = seq_indices < seq_len_expanded.unsqueeze(-1) |
|
mask = mask * (valid_mask.unsqueeze(2) & valid_mask.unsqueeze(3)).to(dtype) |
|
|
|
|
|
if self.causal: |
|
causal_mask = torch.triu(torch.ones(max_len, max_len, device=device, dtype=torch.bool), diagonal=1) |
|
mask = mask * (~causal_mask.unsqueeze(0).unsqueeze(1)).to(dtype) |
|
|
|
|
|
mask = mask + (1.0 - mask) * torch.finfo(dtype).min |
|
return mask |
|
|
|
def forward(self, hidden_states: torch.Tensor, seq_len: torch.Tensor) -> torch.Tensor: |
|
""" |
|
Forward propagation, input and output are [bsz, max_len, embed_dim]. |
|
|
|
Parameters: |
|
hidden_states (torch.Tensor): Input hidden states, shape [bsz, max_len, embed_dim] |
|
seq_len (torch.Tensor): Sequence length for each sample, shape [bsz] |
|
|
|
Returns: |
|
attn_output (torch.Tensor): Attention output, shape [bsz, max_len, embed_dim] |
|
""" |
|
bsz, max_len, _ = hidden_states.size() |
|
|
|
|
|
query = self.q_proj(hidden_states) * self.scaling |
|
key = self.k_proj(hidden_states) |
|
value = self.v_proj(hidden_states) |
|
|
|
|
|
query = query.view(bsz, max_len, self.num_heads, self.head_dim).transpose(1, 2) |
|
key = key.view(bsz, max_len, self.num_heads, self.head_dim).transpose(1, 2) |
|
value = value.view(bsz, max_len, self.num_heads, self.head_dim).transpose(1, 2) |
|
|
|
|
|
attn_scores = torch.matmul(query, key.transpose(-1, -2)) |
|
|
|
|
|
attn_mask = self._create_attention_mask(seq_len, max_len, hidden_states.device, attn_scores.dtype) |
|
|
|
attn_scores = attn_scores + attn_mask |
|
|
|
|
|
attn_weights = F.softmax(attn_scores, dim=-1) |
|
attn_weights = self.dropout(attn_weights) |
|
|
|
|
|
attn_output = torch.matmul(attn_weights, value) |
|
attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, max_len, self.embed_dim) |
|
|
|
|
|
attn_output = self.out_proj(attn_output) |
|
|
|
return attn_output |
|
|
|
|
|
class OmniWhisperTransformerLayer(nn.Module): |
|
def __init__(self, activation_function="gelu", d_model=1280, attention_heads=20, ffn_dim=5120, causal=False, ln_type="LayerNorm", attn_type="varlen"): |
|
super().__init__() |
|
self.embed_dim = d_model |
|
|
|
if attn_type != "varlen": |
|
raise ValueError(f"Unknown attn_type: {attn_type}. Only 'varlen' is supported.") |
|
self.self_attn = VarLenAttention(self.embed_dim, attention_heads, causal) |
|
if ln_type == "LayerNorm": |
|
self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) |
|
elif ln_type == "RMSNorm": |
|
self.self_attn_layer_norm = RMSNorm(self.embed_dim) |
|
else: |
|
raise ValueError(f"Unknown ln_type: {ln_type}") |
|
self.activation_fn = ACT2FN[activation_function] |
|
self.fc1 = nn.Linear(self.embed_dim, ffn_dim) |
|
self.fc2 = nn.Linear(ffn_dim, self.embed_dim) |
|
if ln_type == "LayerNorm": |
|
self.final_layer_norm = nn.LayerNorm(self.embed_dim) |
|
elif ln_type == "RMSNorm": |
|
self.final_layer_norm = RMSNorm(self.embed_dim) |
|
else: |
|
raise ValueError(f"Unknown ln_type: {ln_type}") |
|
|
|
def forward(self, hidden_states: torch.Tensor, seq_len: torch.Tensor) -> torch.Tensor: |
|
residual = hidden_states |
|
hidden_states = self.self_attn_layer_norm(hidden_states) |
|
|
|
|
|
|
|
|
|
hidden_states = self.self_attn(hidden_states, seq_len) |
|
hidden_states = residual + hidden_states |
|
residual = hidden_states |
|
hidden_states = self.final_layer_norm(hidden_states) |
|
hidden_states = self.activation_fn(self.fc1(hidden_states)) |
|
hidden_states = self.fc2(hidden_states) |
|
hidden_states = residual + hidden_states |
|
if (hidden_states.dtype == torch.float16 or hidden_states.dtype == torch.bfloat16) and \ |
|
(torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()): |
|
clamp_value = torch.finfo(hidden_states.dtype).max - 1000 |
|
hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) |
|
return hidden_states |
|
|
|
|
|
class OmniAudioEncoder(nn.Module): |
|
def __init__( |
|
self, |
|
num_mel_bins=128, |
|
sampling_rate=16000, |
|
hop_length=160, |
|
stride_size=2, |
|
kernel_size=3, |
|
d_model=1280, |
|
scale_embedding=True, |
|
max_audio_seconds=30, |
|
encoder_layers=32, |
|
encoder_attention_heads=20, |
|
encoder_ffn_dim=5120, |
|
activation_function="gelu", |
|
attn_type="varlen" |
|
): |
|
super().__init__() |
|
|
|
self.max_source_positions = (max_audio_seconds * sampling_rate // hop_length) // stride_size |
|
|
|
self.embed_scale = math.sqrt(d_model) if scale_embedding else 1.0 |
|
self.num_mel_bins = num_mel_bins |
|
self.d_model = d_model |
|
self.stride_size = stride_size |
|
|
|
|
|
self.conv1 = nn.Conv1d(num_mel_bins, d_model, kernel_size=kernel_size, padding=1) |
|
|
|
self.conv2 = nn.Conv1d(d_model, d_model, kernel_size=kernel_size, stride=stride_size, padding=1) |
|
|
|
|
|
self.register_buffer("positional_embedding", sinusoids(self.max_source_positions, d_model)) |
|
|
|
|
|
self.layers = nn.ModuleList([ |
|
OmniWhisperTransformerLayer( |
|
activation_function=activation_function, |
|
d_model=d_model, |
|
attention_heads=encoder_attention_heads, |
|
ffn_dim=encoder_ffn_dim, |
|
causal=False, |
|
attn_type=attn_type |
|
) for _ in range(encoder_layers) |
|
]) |
|
|
|
|
|
self.layer_norm = nn.LayerNorm(d_model) |
|
|
|
def forward(self, input_features, input_length, output_hidden_states=False): |
|
""" |
|
Forward propagation function to convert input audio features to hidden state representation |
|
|
|
Parameters: |
|
input_features (torch.Tensor): Input Mel spectrogram features, shape [bsz, num_mel_bins, seq_len] |
|
input_length (torch.Tensor): Input sequence length for each sample, shape [bsz] |
|
output_hidden_states (bool, optional): Whether to return hidden states for each layer, default False |
|
|
|
Returns: |
|
if output_hidden_states is False: |
|
hidden_states (torch.Tensor): Encoded hidden states, shape [bsz, d_model, tgt_len] |
|
output_length (torch.Tensor): Output sequence length for each sample, shape [bsz] |
|
else: |
|
hidden_states (torch.Tensor): Encoded hidden states, shape [bsz, d_model, tgt_len] |
|
output_length (torch.Tensor): Output sequence length for each sample, shape [bsz] |
|
hidden_states_all_layers (tuple): Tuple containing hidden states for each layer, including initial input |
|
""" |
|
|
|
input_features = input_features.to(self.conv1.weight.dtype) |
|
|
|
|
|
inputs_embeds = nn.functional.gelu(self.conv1(input_features)) |
|
|
|
|
|
inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds)) |
|
|
|
|
|
output_length = (input_length // self.stride_size).long() |
|
|
|
|
|
hidden_states = inputs_embeds.permute(0, 2, 1) |
|
|
|
|
|
bsz, tgt_len, _ = hidden_states.size() |
|
|
|
|
|
if tgt_len < self.positional_embedding.shape[0]: |
|
current_positional_embedding = self.positional_embedding[:tgt_len] |
|
else: |
|
current_positional_embedding = self.positional_embedding |
|
|
|
|
|
hidden_states = (hidden_states.to(torch.float32) + current_positional_embedding).to(hidden_states.dtype) |
|
|
|
|
|
attention_mask = get_sequence_mask(hidden_states, output_length) |
|
|
|
|
|
hidden_states_all_layers = () if output_hidden_states else None |
|
|
|
|
|
for encoder_layer in self.layers: |
|
if output_hidden_states: |
|
hidden_states_all_layers = hidden_states_all_layers + (hidden_states,) |
|
hidden_states = encoder_layer(hidden_states, output_length) |
|
|
|
|
|
hidden_states = self.layer_norm(hidden_states) |
|
if output_hidden_states: |
|
hidden_states_all_layers = hidden_states_all_layers + (hidden_states,) |
|
|
|
|
|
hidden_states = torch.where(attention_mask, hidden_states, 0) |
|
hidden_states = hidden_states.transpose(1, 2) |
|
|
|
if not output_hidden_states: |
|
return hidden_states, output_length |
|
else: |
|
return hidden_states, output_length, hidden_states_all_layers |
|
|
|
|
|
class OmniAudioDecoder(nn.Module): |
|
def __init__( |
|
self, |
|
num_mel_bins=128, |
|
sampling_rate=16000, |
|
hop_length=160, |
|
stride_size=2, |
|
kernel_size=3, |
|
d_model=1280, |
|
scale_embedding=True, |
|
max_audio_seconds=30, |
|
decoder_layers=32, |
|
decoder_attention_heads=20, |
|
decoder_ffn_dim=5120, |
|
activation_function="gelu", |
|
attn_type="varlen" |
|
): |
|
super().__init__() |
|
self.max_source_positions = (max_audio_seconds * sampling_rate // hop_length) // stride_size |
|
self.embed_scale = math.sqrt(d_model) if scale_embedding else 1.0 |
|
self.num_mel_bins = num_mel_bins |
|
self.d_model = d_model |
|
self.stride_size = stride_size |
|
|
|
|
|
self.deconv1 = nn.ConvTranspose1d( |
|
d_model, |
|
d_model, |
|
kernel_size=kernel_size, |
|
stride=stride_size, |
|
padding=0, |
|
output_padding=0 |
|
) |
|
self.deconv2 = nn.ConvTranspose1d( |
|
d_model, |
|
num_mel_bins, |
|
kernel_size=kernel_size, |
|
stride=1, |
|
padding=0 |
|
) |
|
|
|
|
|
self.register_buffer("positional_embedding", sinusoids(self.max_source_positions, d_model)) |
|
|
|
|
|
self.layers = nn.ModuleList([ |
|
OmniWhisperTransformerLayer( |
|
activation_function=activation_function, |
|
d_model=d_model, |
|
attention_heads=decoder_attention_heads, |
|
ffn_dim=decoder_ffn_dim, |
|
causal=False, |
|
attn_type=attn_type |
|
) for _ in range(decoder_layers) |
|
]) |
|
self.layer_norm = nn.LayerNorm(d_model) |
|
|
|
def forward(self, hidden_states, input_length): |
|
|
|
hidden_states = hidden_states.transpose(1, 2) |
|
bsz, tgt_len, _ = hidden_states.size() |
|
|
|
|
|
if tgt_len < self.positional_embedding.shape[0]: |
|
current_positional_embedding = self.positional_embedding[:tgt_len] |
|
else: |
|
current_positional_embedding = self.positional_embedding |
|
hidden_states = (hidden_states.to(torch.float32) + current_positional_embedding).to(hidden_states.dtype) |
|
|
|
|
|
attention_mask = get_sequence_mask(hidden_states, input_length) |
|
|
|
|
|
for decoder_layer in self.layers: |
|
hidden_states = decoder_layer(hidden_states, input_length) |
|
|
|
|
|
hidden_states = self.layer_norm(hidden_states) |
|
|
|
|
|
hidden_states = torch.where(attention_mask, hidden_states, 0) |
|
|
|
|
|
hidden_states = hidden_states.permute(0, 2, 1) |
|
output_features = nn.functional.gelu(self.deconv1(hidden_states)) |
|
output_features = nn.functional.gelu(self.deconv2(output_features)) |
|
|
|
|
|
expected_length = tgt_len * self.stride_size |
|
if output_features.size(2) > expected_length: |
|
output_features = output_features[:, :, :expected_length] |
|
|
|
output_length = input_length * self.stride_size |
|
|
|
return output_features, output_length |
|
|
|
|
|
class ResidualDownConv(nn.Module): |
|
def __init__(self, d_model=1280, avg_pooler=4): |
|
""" |
|
Downsampling module containing residual connection and convolution operation |
|
|
|
Parameters: |
|
d_model (int): Input and output hidden dimension |
|
avg_pooler (int): Downsampling factor (convolution step) |
|
""" |
|
super().__init__() |
|
self.d_model = d_model |
|
self.avg_pooler = avg_pooler |
|
self.intermediate_dim = d_model * avg_pooler |
|
|
|
|
|
self.gate_proj = nn.Conv1d(d_model, self.intermediate_dim, avg_pooler, avg_pooler, bias=False) |
|
self.up_proj = nn.Conv1d(d_model, self.intermediate_dim, avg_pooler, avg_pooler, bias=False) |
|
|
|
|
|
self.down_proj = nn.Linear(self.intermediate_dim, self.intermediate_dim, bias=False) |
|
|
|
|
|
self.act_fn = ACT2FN['silu'] |
|
self.layer_norm = nn.LayerNorm(self.intermediate_dim) |
|
|
|
def forward(self, x, input_length): |
|
""" |
|
Forward propagation, execute downsampling and residual processing |
|
|
|
Parameters: |
|
x (torch.Tensor): Input tensor, shape [B, D, T] |
|
|
|
Returns: |
|
res (torch.Tensor): Downsampled feature, shape [B, intermediate_dim, seq_len // avg_pooler] |
|
valid_mask (torch.Tensor): Valid sequence mask |
|
""" |
|
output_length = input_length // self.avg_pooler |
|
x = x.transpose(1, 2) |
|
batch_size, seq_len, _ = x.shape |
|
if seq_len % self.avg_pooler != 0: |
|
pad_size = self.avg_pooler - seq_len % self.avg_pooler |
|
x = F.pad(x, (0, pad_size), "constant", 0) |
|
|
|
xt = x.permute(0, 2, 1) |
|
g = self.gate_proj(xt).permute(0, 2, 1) |
|
u = self.up_proj(xt).permute(0, 2, 1) |
|
x = x.reshape(batch_size, -1, self.intermediate_dim) |
|
|
|
c = self.down_proj(self.act_fn(g) * u) |
|
res = self.layer_norm(c + x) |
|
res = res.transpose(1, 2) |
|
return res, output_length |
|
|
|
|
|
class UpConv(nn.Module): |
|
def __init__(self, d_model=1280, stride=4): |
|
""" |
|
Simple upsampling module using transpose convolution |
|
|
|
Parameters: |
|
d_model (int): Input and output hidden dimension |
|
stride (int): Upsampling factor (transpose convolution step) |
|
""" |
|
super().__init__() |
|
self.d_model = d_model |
|
self.stride = stride |
|
|
|
|
|
self.up_conv = nn.ConvTranspose1d( |
|
self.stride * d_model, |
|
d_model, |
|
kernel_size=stride, |
|
stride=stride, |
|
bias=False |
|
) |
|
|
|
def forward(self, x, input_length): |
|
""" |
|
Forward propagation, execute upsampling |
|
|
|
Parameters: |
|
x (torch.Tensor): Input tensor, shape [B, D * stride, T] |
|
|
|
Returns: |
|
res (torch.Tensor): Upsampled feature, shape [B, D, T * stride] |
|
""" |
|
|
|
res = self.up_conv(x) |
|
output_length = input_length * self.stride |
|
return res, output_length |
|
|
|
|
|
|
|
class Transformer(nn.Module): |
|
def __init__( |
|
self, |
|
input_dim=1280, |
|
d_model=1280, |
|
output_dim=1280, |
|
max_source_positions=1500, |
|
encoder_layers=32, |
|
encoder_attention_heads=20, |
|
encoder_ffn_dim=5120, |
|
activation_function="gelu", |
|
attn_type="varlen" |
|
): |
|
super().__init__() |
|
self.input_dim = input_dim |
|
self.d_model = d_model |
|
self.output_dim = output_dim |
|
self.max_source_positions = max_source_positions |
|
|
|
|
|
if input_dim != d_model: |
|
self.proj = nn.Linear(input_dim, d_model, bias=True) |
|
else: |
|
self.proj = None |
|
|
|
|
|
self.register_buffer("positional_embedding", sinusoids(self.max_source_positions, d_model)) |
|
|
|
|
|
self.layers = nn.ModuleList([ |
|
OmniWhisperTransformerLayer( |
|
activation_function=activation_function, |
|
d_model=d_model, |
|
attention_heads=encoder_attention_heads, |
|
ffn_dim=encoder_ffn_dim, |
|
causal=False, |
|
attn_type=attn_type |
|
) for _ in range(encoder_layers) |
|
]) |
|
|
|
|
|
self.layer_norm = nn.LayerNorm(d_model) |
|
|
|
|
|
if output_dim != d_model: |
|
self.out_proj = nn.Linear(d_model, output_dim, bias=True) |
|
else: |
|
self.out_proj = None |
|
|
|
def forward(self, input_features: torch.Tensor, input_length: torch.Tensor, output_hidden_states: bool = False): |
|
""" |
|
Forward propagation function to convert input features through Transformer layer to hidden state representation |
|
|
|
Parameters: |
|
input_features (torch.Tensor): Input features, shape [bsz, input_dim, seq_len] (B, input_dim, T) |
|
input_length (torch.Tensor): Input sequence length for each sample, shape [bsz] |
|
output_hidden_states (bool, optional): Whether to return hidden states for each layer, default False |
|
|
|
Returns: |
|
if output_hidden_states is False: |
|
hidden_states (torch.Tensor): Encoded hidden states, shape [bsz, output_dim, seq_len] (B, output_dim, T) |
|
output_length (torch.Tensor): Output sequence length for each sample, shape [bsz] |
|
else: |
|
hidden_states (torch.Tensor): Encoded hidden states, shape [bsz, output_dim, seq_len] (B, output_dim, T) |
|
output_length (torch.Tensor): Output sequence length for each sample, shape [bsz] |
|
hidden_states_all_layers (tuple): Tuple containing hidden states for each layer, each shape [bsz, seq_len, d_model] |
|
""" |
|
|
|
output_length = input_length.long() |
|
|
|
|
|
if self.proj is not None: |
|
hidden_states = self.proj(input_features.permute(0, 2, 1)).permute(0, 2, 1) |
|
else: |
|
hidden_states = input_features |
|
|
|
|
|
hidden_states = hidden_states.permute(0, 2, 1) |
|
|
|
|
|
bsz, tgt_len, _ = hidden_states.size() |
|
|
|
|
|
if tgt_len < self.positional_embedding.shape[0]: |
|
current_positional_embedding = self.positional_embedding[:tgt_len] |
|
else: |
|
current_positional_embedding = self.positional_embedding |
|
|
|
|
|
hidden_states = (hidden_states.to(torch.float32) + current_positional_embedding).to(hidden_states.dtype) |
|
|
|
|
|
attention_mask = get_sequence_mask(hidden_states, output_length) |
|
|
|
|
|
hidden_states_all_layers = () if output_hidden_states else None |
|
|
|
|
|
for encoder_layer in self.layers: |
|
if output_hidden_states: |
|
hidden_states_all_layers = hidden_states_all_layers + (hidden_states,) |
|
hidden_states = encoder_layer(hidden_states, output_length) |
|
|
|
|
|
hidden_states = self.layer_norm(hidden_states) |
|
if output_hidden_states: |
|
hidden_states_all_layers = hidden_states_all_layers + (hidden_states,) |
|
|
|
|
|
hidden_states = torch.where(attention_mask, hidden_states, 0) |
|
|
|
|
|
hidden_states = hidden_states.transpose(1, 2) |
|
|
|
|
|
if self.out_proj is not None: |
|
hidden_states = self.out_proj(hidden_states.permute(0, 2, 1)).permute(0, 2, 1) |
|
|
|
if not output_hidden_states: |
|
return hidden_states, output_length |
|
else: |
|
return hidden_states, output_length, hidden_states_all_layers |
|
|
|
|
|
def safe_log(x: torch.Tensor, clip_val: float = 1e-7) -> torch.Tensor: |
|
""" |
|
Computes the element-wise logarithm of the input tensor with clipping to avoid near-zero values. |
|
|
|
Args: |
|
x (Tensor): Input tensor. |
|
clip_val (float, optional): Minimum value to clip the input tensor. Defaults to 1e-7. |
|
|
|
Returns: |
|
Tensor: Element-wise logarithm of the input tensor with clipping applied. |
|
""" |
|
return torch.log(torch.clip(x, min=clip_val)) |
|
|
|
|
|
def symlog(x: torch.Tensor) -> torch.Tensor: |
|
return torch.sign(x) * torch.log1p(x.abs()) |
|
|
|
|
|
def symexp(x: torch.Tensor) -> torch.Tensor: |
|
return torch.sign(x) * (torch.exp(x.abs()) - 1) |
|
|
|
|
|
class STFT(nn.Module): |
|
def __init__( |
|
self, |
|
n_fft: int, |
|
hop_length: int, |
|
win_length: int, |
|
center=True, |
|
): |
|
super().__init__() |
|
self.center = center |
|
self.n_fft = n_fft |
|
self.hop_length = hop_length |
|
self.win_length = win_length |
|
window = torch.hann_window(win_length) |
|
self.register_buffer("window", window) |
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
|
|
|
|
if not self.center: |
|
pad = self.win_length - self.hop_length |
|
x = torch.nn.functional.pad(x, (pad // 2, pad // 2), mode="reflect") |
|
|
|
stft_spec = torch.stft( |
|
x, |
|
self.n_fft, |
|
hop_length=self.hop_length, |
|
win_length=self.win_length, |
|
window=self.window, |
|
center=self.center, |
|
return_complex=False, |
|
) |
|
|
|
rea = stft_spec[:, :, :, 0] |
|
imag = stft_spec[:, :, :, 1] |
|
|
|
log_mag = torch.log( |
|
torch.abs(torch.sqrt(torch.pow(rea, 2) + torch.pow(imag, 2))) + 1e-5 |
|
) |
|
phase = torch.atan2(imag, rea) |
|
|
|
return log_mag, phase |
|
|
|
|
|
class ISTFT(nn.Module): |
|
""" |
|
Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with |
|
windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges. |
|
See issue: https://github.com/pytorch/pytorch/issues/62323 |
|
Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs. |
|
The NOLA constraint is met as we trim padded samples anyway. |
|
|
|
Args: |
|
n_fft (int): Size of Fourier transform. |
|
hop_length (int): The distance between neighboring sliding window frames. |
|
win_length (int): The size of window frame and STFT filter. |
|
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". |
|
""" |
|
|
|
def __init__( |
|
self, n_fft: int, hop_length: int, win_length: int, padding: str = "same" |
|
): |
|
super().__init__() |
|
if padding not in ["center", "same"]: |
|
raise ValueError("Padding must be 'center' or 'same'.") |
|
self.padding = padding |
|
self.n_fft = n_fft |
|
self.hop_length = hop_length |
|
self.win_length = win_length |
|
window = torch.hann_window(win_length) |
|
self.register_buffer("window", window) |
|
|
|
def forward(self, spec: torch.Tensor) -> torch.Tensor: |
|
""" |
|
Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram. |
|
|
|
Args: |
|
spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size, |
|
N is the number of frequency bins, and T is the number of time frames. |
|
|
|
Returns: |
|
Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal. |
|
""" |
|
if self.padding == "center": |
|
|
|
return torch.istft( |
|
spec, |
|
self.n_fft, |
|
self.hop_length, |
|
self.win_length, |
|
self.window, |
|
center=True, |
|
) |
|
elif self.padding == "same": |
|
pad = (self.win_length - self.hop_length) // 2 |
|
else: |
|
raise ValueError("Padding must be 'center' or 'same'.") |
|
|
|
assert spec.dim() == 3, "Expected a 3D tensor as input" |
|
B, N, T = spec.shape |
|
|
|
|
|
ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") |
|
ifft = ifft * self.window[None, :, None] |
|
|
|
|
|
output_size = (T - 1) * self.hop_length + self.win_length |
|
y = torch.nn.functional.fold( |
|
ifft, |
|
output_size=(1, output_size), |
|
kernel_size=(1, self.win_length), |
|
stride=(1, self.hop_length), |
|
)[:, 0, 0, pad:-pad] |
|
|
|
|
|
window_sq = self.window.square().expand(1, T, -1).transpose(1, 2) |
|
window_envelope = torch.nn.functional.fold( |
|
window_sq, |
|
output_size=(1, output_size), |
|
kernel_size=(1, self.win_length), |
|
stride=(1, self.hop_length), |
|
).squeeze()[pad:-pad] |
|
|
|
|
|
assert (window_envelope > 1e-11).all() |
|
y = y / window_envelope |
|
|
|
return y |
|
|
|
|
|
class MDCT(nn.Module): |
|
""" |
|
Modified Discrete Cosine Transform (MDCT) module. |
|
|
|
Args: |
|
frame_len (int): Length of the MDCT frame. |
|
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". |
|
""" |
|
|
|
def __init__(self, frame_len: int, padding: str = "same"): |
|
super().__init__() |
|
if padding not in ["center", "same"]: |
|
raise ValueError("Padding must be 'center' or 'same'.") |
|
self.padding = padding |
|
self.frame_len = frame_len |
|
N = frame_len // 2 |
|
n0 = (N + 1) / 2 |
|
window = torch.from_numpy(scipy.signal.cosine(frame_len)).float() |
|
self.register_buffer("window", window) |
|
|
|
pre_twiddle = torch.exp(-1j * torch.pi * torch.arange(frame_len) / frame_len) |
|
post_twiddle = torch.exp(-1j * torch.pi * n0 * (torch.arange(N) + 0.5) / N) |
|
|
|
|
|
self.register_buffer("pre_twiddle", view_as_real(pre_twiddle)) |
|
self.register_buffer("post_twiddle", view_as_real(post_twiddle)) |
|
|
|
def forward(self, audio: torch.Tensor) -> torch.Tensor: |
|
""" |
|
Apply the Modified Discrete Cosine Transform (MDCT) to the input audio. |
|
|
|
Args: |
|
audio (Tensor): Input audio waveform of shape (B, T), where B is the batch size |
|
and T is the length of the audio. |
|
|
|
Returns: |
|
Tensor: MDCT coefficients of shape (B, L, N), where L is the number of output frames |
|
and N is the number of frequency bins. |
|
""" |
|
if self.padding == "center": |
|
audio = torch.nn.functional.pad( |
|
audio, (self.frame_len // 2, self.frame_len // 2) |
|
) |
|
elif self.padding == "same": |
|
|
|
audio = torch.nn.functional.pad( |
|
audio, (self.frame_len // 4, self.frame_len // 4) |
|
) |
|
else: |
|
raise ValueError("Padding must be 'center' or 'same'.") |
|
|
|
x = audio.unfold(-1, self.frame_len, self.frame_len // 2) |
|
N = self.frame_len // 2 |
|
x = x * self.window.expand(x.shape) |
|
X = torch.fft.fft( |
|
x * view_as_complex(self.pre_twiddle).expand(x.shape), dim=-1 |
|
)[..., :N] |
|
res = X * view_as_complex(self.post_twiddle).expand(X.shape) * np.sqrt(1 / N) |
|
return torch.real(res) * np.sqrt(2) |
|
|
|
|
|
class IMDCT(nn.Module): |
|
""" |
|
Inverse Modified Discrete Cosine Transform (IMDCT) module. |
|
|
|
Args: |
|
frame_len (int): Length of the MDCT frame. |
|
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". |
|
""" |
|
|
|
def __init__(self, frame_len: int, padding: str = "same"): |
|
super().__init__() |
|
if padding not in ["center", "same"]: |
|
raise ValueError("Padding must be 'center' or 'same'.") |
|
self.padding = padding |
|
self.frame_len = frame_len |
|
N = frame_len // 2 |
|
n0 = (N + 1) / 2 |
|
window = torch.from_numpy(scipy.signal.cosine(frame_len)).float() |
|
self.register_buffer("window", window) |
|
|
|
pre_twiddle = torch.exp(1j * torch.pi * n0 * torch.arange(N * 2) / N) |
|
post_twiddle = torch.exp(1j * torch.pi * (torch.arange(N * 2) + n0) / (N * 2)) |
|
self.register_buffer("pre_twiddle", view_as_real(pre_twiddle)) |
|
self.register_buffer("post_twiddle", view_as_real(post_twiddle)) |
|
|
|
def forward(self, X: torch.Tensor) -> torch.Tensor: |
|
""" |
|
Apply the Inverse Modified Discrete Cosine Transform (IMDCT) to the input MDCT coefficients. |
|
|
|
Args: |
|
X (Tensor): Input MDCT coefficients of shape (B, L, N), where B is the batch size, |
|
L is the number of frames, and N is the number of frequency bins. |
|
|
|
Returns: |
|
Tensor: Reconstructed audio waveform of shape (B, T), where T is the length of the audio. |
|
""" |
|
B, L, N = X.shape |
|
Y = torch.zeros((B, L, N * 2), dtype=X.dtype, device=X.device) |
|
Y[..., :N] = X |
|
Y[..., N:] = -1 * torch.conj(torch.flip(X, dims=(-1,))) |
|
y = torch.fft.ifft( |
|
Y * view_as_complex(self.pre_twiddle).expand(Y.shape), dim=-1 |
|
) |
|
y = ( |
|
torch.real(y * view_as_complex(self.post_twiddle).expand(y.shape)) |
|
* np.sqrt(N) |
|
* np.sqrt(2) |
|
) |
|
result = y * self.window.expand(y.shape) |
|
output_size = (1, (L + 1) * N) |
|
audio = torch.nn.functional.fold( |
|
result.transpose(1, 2), |
|
output_size=output_size, |
|
kernel_size=(1, self.frame_len), |
|
stride=(1, self.frame_len // 2), |
|
)[:, 0, 0, :] |
|
|
|
if self.padding == "center": |
|
pad = self.frame_len // 2 |
|
elif self.padding == "same": |
|
pad = self.frame_len // 4 |
|
else: |
|
raise ValueError("Padding must be 'center' or 'same'.") |
|
|
|
audio = audio[:, pad:-pad] |
|
return audio |
|
|
|
|
|
class FourierHead(nn.Module): |
|
"""Base class for inverse fourier modules.""" |
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
""" |
|
Args: |
|
x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, |
|
L is the sequence length, and H denotes the model dimension. |
|
|
|
Returns: |
|
Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. |
|
""" |
|
raise NotImplementedError("Subclasses must implement the forward method.") |
|
|
|
|
|
class ISTFTHead(FourierHead): |
|
""" |
|
ISTFT Head module for predicting STFT complex coefficients. |
|
|
|
Args: |
|
dim (int): Hidden dimension of the model. |
|
n_fft (int): Size of Fourier transform. |
|
hop_length (int): The distance between neighboring sliding window frames, which should align with |
|
the resolution of the input features. |
|
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". |
|
""" |
|
|
|
def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"): |
|
super().__init__() |
|
out_dim = n_fft + 2 |
|
self.out = torch.nn.Linear(dim, out_dim) |
|
self.istft = ISTFT( |
|
n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding |
|
) |
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
""" |
|
Forward pass of the ISTFTHead module. |
|
|
|
Args: |
|
x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, |
|
L is the sequence length, and H denotes the model dimension. |
|
|
|
Returns: |
|
Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. |
|
""" |
|
x = self.out(x).transpose(1, 2) |
|
mag, p = x.chunk(2, dim=1) |
|
mag = torch.exp(mag) |
|
mag = torch.clip( |
|
mag, max=1e2 |
|
) |
|
|
|
x = torch.cos(p) |
|
y = torch.sin(p) |
|
|
|
|
|
|
|
|
|
|
|
original_dtype = x.dtype |
|
S = mag.float() * (x.float() + 1j * y.float()) |
|
audio = self.istft(S) |
|
audio = audio.to(original_dtype) |
|
return audio |
|
|
|
|
|
class IMDCTSymExpHead(FourierHead): |
|
""" |
|
IMDCT Head module for predicting MDCT coefficients with symmetric exponential function |
|
|
|
Args: |
|
dim (int): Hidden dimension of the model. |
|
mdct_frame_len (int): Length of the MDCT frame. |
|
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". |
|
sample_rate (int, optional): The sample rate of the audio. If provided, the last layer will be initialized |
|
based on perceptual scaling. Defaults to None. |
|
clip_audio (bool, optional): Whether to clip the audio output within the range of [-1.0, 1.0]. Defaults to False. |
|
""" |
|
|
|
def __init__( |
|
self, |
|
dim: int, |
|
mdct_frame_len: int, |
|
padding: str = "same", |
|
sample_rate: Optional[int] = None, |
|
clip_audio: bool = False, |
|
): |
|
super().__init__() |
|
out_dim = mdct_frame_len // 2 |
|
self.out = nn.Linear(dim, out_dim) |
|
self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding) |
|
self.clip_audio = clip_audio |
|
|
|
if sample_rate is not None: |
|
|
|
m_max = _hz_to_mel(sample_rate // 2) |
|
m_pts = torch.linspace(0, m_max, out_dim) |
|
f_pts = _mel_to_hz(m_pts) |
|
scale = 1 - (f_pts / f_pts.max()) |
|
|
|
with torch.no_grad(): |
|
self.out.weight.mul_(scale.view(-1, 1)) |
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
""" |
|
Forward pass of the IMDCTSymExpHead module. |
|
|
|
Args: |
|
x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, |
|
L is the sequence length, and H denotes the model dimension. |
|
|
|
Returns: |
|
Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. |
|
""" |
|
x = self.out(x) |
|
x = symexp(x) |
|
x = torch.clip( |
|
x, min=-1e2, max=1e2 |
|
) |
|
audio = self.imdct(x) |
|
if self.clip_audio: |
|
audio = torch.clip(x, min=-1.0, max=1.0) |
|
|
|
return audio |
|
|
|
|
|
class IMDCTCosHead(FourierHead): |
|
""" |
|
IMDCT Head module for predicting MDCT coefficients with parametrizing MDCT = exp(m) · cos(p) |
|
|
|
Args: |
|
dim (int): Hidden dimension of the model. |
|
mdct_frame_len (int): Length of the MDCT frame. |
|
padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". |
|
clip_audio (bool, optional): Whether to clip the audio output within the range of [-1.0, 1.0]. Defaults to False. |
|
""" |
|
|
|
def __init__( |
|
self, |
|
dim: int, |
|
mdct_frame_len: int, |
|
padding: str = "same", |
|
clip_audio: bool = False, |
|
): |
|
super().__init__() |
|
self.clip_audio = clip_audio |
|
self.out = nn.Linear(dim, mdct_frame_len) |
|
self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding) |
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
""" |
|
Forward pass of the IMDCTCosHead module. |
|
|
|
Args: |
|
x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, |
|
L is the sequence length, and H denotes the model dimension. |
|
|
|
Returns: |
|
Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. |
|
""" |
|
x = self.out(x) |
|
m, p = x.chunk(2, dim=2) |
|
m = torch.exp(m).clip( |
|
max=1e2 |
|
) |
|
audio = self.imdct(m * torch.cos(p)) |
|
if self.clip_audio: |
|
audio = torch.clip(x, min=-1.0, max=1.0) |
|
return audio |
|
|
|
|
|
class ConvNeXtBlock(nn.Module): |
|
"""ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal. |
|
|
|
Args: |
|
dim (int): Number of input channels. |
|
intermediate_dim (int): Dimensionality of the intermediate layer. |
|
layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling. |
|
Defaults to None. |
|
adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm. |
|
None means non-conditional LayerNorm. Defaults to None. |
|
""" |
|
|
|
def __init__( |
|
self, |
|
dim: int, |
|
intermediate_dim: int, |
|
layer_scale_init_value: float, |
|
adanorm_num_embeddings: Optional[int] = None, |
|
): |
|
super().__init__() |
|
self.dwconv = nn.Conv1d( |
|
dim, dim, kernel_size=7, padding=3, groups=dim |
|
) |
|
self.adanorm = adanorm_num_embeddings is not None |
|
if adanorm_num_embeddings: |
|
self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6) |
|
else: |
|
self.norm = nn.LayerNorm(dim, eps=1e-6) |
|
self.pwconv1 = nn.Linear( |
|
dim, intermediate_dim |
|
) |
|
self.act = nn.GELU() |
|
self.pwconv2 = nn.Linear(intermediate_dim, dim) |
|
self.gamma = ( |
|
nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True) |
|
if layer_scale_init_value > 0 |
|
else None |
|
) |
|
|
|
def forward( |
|
self, x: torch.Tensor, cond_embedding_id: Optional[torch.Tensor] = None |
|
) -> torch.Tensor: |
|
residual = x |
|
x = self.dwconv(x) |
|
x = x.transpose(1, 2) |
|
if self.adanorm: |
|
assert cond_embedding_id is not None |
|
x = self.norm(x, cond_embedding_id) |
|
else: |
|
x = self.norm(x) |
|
x = self.pwconv1(x) |
|
x = self.act(x) |
|
x = self.pwconv2(x) |
|
if self.gamma is not None: |
|
x = self.gamma * x |
|
x = x.transpose(1, 2) |
|
|
|
x = residual + x |
|
return x |
|
|
|
|
|
class AdaLayerNorm(nn.Module): |
|
""" |
|
Adaptive Layer Normalization module with learnable embeddings per `num_embeddings` classes |
|
|
|
Args: |
|
num_embeddings (int): Number of embeddings. |
|
embedding_dim (int): Dimension of the embeddings. |
|
""" |
|
|
|
def __init__(self, num_embeddings: int, embedding_dim: int, eps: float = 1e-6): |
|
super().__init__() |
|
self.eps = eps |
|
self.dim = embedding_dim |
|
self.scale = nn.Embedding( |
|
num_embeddings=num_embeddings, embedding_dim=embedding_dim |
|
) |
|
self.shift = nn.Embedding( |
|
num_embeddings=num_embeddings, embedding_dim=embedding_dim |
|
) |
|
torch.nn.init.ones_(self.scale.weight) |
|
torch.nn.init.zeros_(self.shift.weight) |
|
|
|
def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor) -> torch.Tensor: |
|
scale = self.scale(cond_embedding_id) |
|
shift = self.shift(cond_embedding_id) |
|
x = nn.functional.layer_norm(x, (self.dim,), eps=self.eps) |
|
x = x * scale + shift |
|
return x |
|
|
|
|
|
class ResBlock1(nn.Module): |
|
""" |
|
ResBlock adapted from HiFi-GAN V1 (https://github.com/jik876/hifi-gan) with dilated 1D convolutions, |
|
but without upsampling layers. |
|
|
|
Args: |
|
dim (int): Number of input channels. |
|
kernel_size (int, optional): Size of the convolutional kernel. Defaults to 3. |
|
dilation (tuple[int], optional): Dilation factors for the dilated convolutions. |
|
Defaults to (1, 3, 5). |
|
lrelu_slope (float, optional): Negative slope of the LeakyReLU activation function. |
|
Defaults to 0.1. |
|
layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling. |
|
Defaults to None. |
|
""" |
|
|
|
def __init__( |
|
self, |
|
dim: int, |
|
kernel_size: int = 3, |
|
dilation: Tuple[int, int, int] = (1, 3, 5), |
|
lrelu_slope: float = 0.1, |
|
layer_scale_init_value: Optional[float] = None, |
|
): |
|
super().__init__() |
|
self.lrelu_slope = lrelu_slope |
|
self.convs1 = nn.ModuleList( |
|
[ |
|
weight_norm( |
|
nn.Conv1d( |
|
dim, |
|
dim, |
|
kernel_size, |
|
1, |
|
dilation=dilation[0], |
|
padding=self.get_padding(kernel_size, dilation[0]), |
|
) |
|
), |
|
weight_norm( |
|
nn.Conv1d( |
|
dim, |
|
dim, |
|
kernel_size, |
|
1, |
|
dilation=dilation[1], |
|
padding=self.get_padding(kernel_size, dilation[1]), |
|
) |
|
), |
|
weight_norm( |
|
nn.Conv1d( |
|
dim, |
|
dim, |
|
kernel_size, |
|
1, |
|
dilation=dilation[2], |
|
padding=self.get_padding(kernel_size, dilation[2]), |
|
) |
|
), |
|
] |
|
) |
|
|
|
self.convs2 = nn.ModuleList( |
|
[ |
|
weight_norm( |
|
nn.Conv1d( |
|
dim, |
|
dim, |
|
kernel_size, |
|
1, |
|
dilation=1, |
|
padding=self.get_padding(kernel_size, 1), |
|
) |
|
), |
|
weight_norm( |
|
nn.Conv1d( |
|
dim, |
|
dim, |
|
kernel_size, |
|
1, |
|
dilation=1, |
|
padding=self.get_padding(kernel_size, 1), |
|
) |
|
), |
|
weight_norm( |
|
nn.Conv1d( |
|
dim, |
|
dim, |
|
kernel_size, |
|
1, |
|
dilation=1, |
|
padding=self.get_padding(kernel_size, 1), |
|
) |
|
), |
|
] |
|
) |
|
|
|
self.gamma = nn.ParameterList( |
|
[ |
|
( |
|
nn.Parameter( |
|
layer_scale_init_value * torch.ones(dim, 1), requires_grad=True |
|
) |
|
if layer_scale_init_value is not None |
|
else None |
|
), |
|
( |
|
nn.Parameter( |
|
layer_scale_init_value * torch.ones(dim, 1), requires_grad=True |
|
) |
|
if layer_scale_init_value is not None |
|
else None |
|
), |
|
( |
|
nn.Parameter( |
|
layer_scale_init_value * torch.ones(dim, 1), requires_grad=True |
|
) |
|
if layer_scale_init_value is not None |
|
else None |
|
), |
|
] |
|
) |
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
for c1, c2, gamma in zip(self.convs1, self.convs2, self.gamma): |
|
xt = torch.nn.functional.leaky_relu(x, negative_slope=self.lrelu_slope) |
|
xt = c1(xt) |
|
xt = torch.nn.functional.leaky_relu(xt, negative_slope=self.lrelu_slope) |
|
xt = c2(xt) |
|
if gamma is not None: |
|
xt = gamma * xt |
|
x = xt + x |
|
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) |
|
|
|
@staticmethod |
|
def get_padding(kernel_size: int, dilation: int = 1) -> int: |
|
return int((kernel_size * dilation - dilation) / 2) |
|
|
|
|
|
class Backbone(nn.Module): |
|
"""Base class for the generator's backbone. It preserves the same temporal resolution across all layers.""" |
|
|
|
def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: |
|
""" |
|
Args: |
|
x (Tensor): Input tensor of shape (B, C, L), where B is the batch size, |
|
C denotes output features, and L is the sequence length. |
|
|
|
Returns: |
|
Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length, |
|
and H denotes the model dimension. |
|
""" |
|
raise NotImplementedError("Subclasses must implement the forward method.") |
|
|
|
|
|
class VocosBackbone(Backbone): |
|
""" |
|
Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization |
|
|
|
Args: |
|
input_channels (int): Number of input features channels. |
|
dim (int): Hidden dimension of the model. |
|
intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock. |
|
num_layers (int): Number of ConvNeXtBlock layers. |
|
layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`. |
|
adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm. |
|
None means non-conditional model. Defaults to None. |
|
""" |
|
|
|
def __init__( |
|
self, |
|
input_channels: int, |
|
dim: int, |
|
intermediate_dim: int, |
|
num_layers: int, |
|
layer_scale_init_value: Optional[float] = None, |
|
adanorm_num_embeddings: Optional[int] = None, |
|
): |
|
super().__init__() |
|
self.input_channels = input_channels |
|
self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3) |
|
self.adanorm = adanorm_num_embeddings is not None |
|
if adanorm_num_embeddings: |
|
self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6) |
|
else: |
|
self.norm = nn.LayerNorm(dim, eps=1e-6) |
|
layer_scale_init_value = layer_scale_init_value or 1 / num_layers |
|
self.convnext = nn.ModuleList( |
|
[ |
|
ConvNeXtBlock( |
|
dim=dim, |
|
intermediate_dim=intermediate_dim, |
|
layer_scale_init_value=layer_scale_init_value, |
|
adanorm_num_embeddings=adanorm_num_embeddings, |
|
) |
|
for _ in range(num_layers) |
|
] |
|
) |
|
self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6) |
|
self.apply(self._init_weights) |
|
|
|
def _init_weights(self, m): |
|
if isinstance(m, (nn.Conv1d, nn.Linear)): |
|
nn.init.trunc_normal_(m.weight, std=0.02) |
|
nn.init.constant_(m.bias, 0) |
|
|
|
def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: |
|
bandwidth_id = kwargs.get("bandwidth_id", None) |
|
x = self.embed(x) |
|
if self.adanorm: |
|
assert bandwidth_id is not None |
|
x = self.norm(x.transpose(1, 2), cond_embedding_id=bandwidth_id) |
|
else: |
|
x = self.norm(x.transpose(1, 2)) |
|
x = x.transpose(1, 2) |
|
for conv_block in self.convnext: |
|
x = conv_block(x, cond_embedding_id=bandwidth_id) |
|
x = self.final_layer_norm(x.transpose(1, 2)) |
|
return x |
|
|
|
|
|
class VocosResNetBackbone(Backbone): |
|
""" |
|
Vocos backbone module built with ResBlocks. |
|
|
|
Args: |
|
input_channels (int): Number of input features channels. |
|
dim (int): Hidden dimension of the model. |
|
num_blocks (int): Number of ResBlock1 blocks. |
|
layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to None. |
|
""" |
|
|
|
def __init__( |
|
self, |
|
input_channels, |
|
dim, |
|
num_blocks, |
|
layer_scale_init_value=None, |
|
): |
|
super().__init__() |
|
self.input_channels = input_channels |
|
self.embed = weight_norm( |
|
nn.Conv1d(input_channels, dim, kernel_size=3, padding=1) |
|
) |
|
layer_scale_init_value = layer_scale_init_value or 1 / num_blocks / 3 |
|
self.resnet = nn.Sequential( |
|
*[ |
|
ResBlock1(dim=dim, layer_scale_init_value=layer_scale_init_value) |
|
for _ in range(num_blocks) |
|
] |
|
) |
|
|
|
def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: |
|
x = self.embed(x) |
|
x = self.resnet(x) |
|
x = x.transpose(1, 2) |
|
return x |
|
|
|
|
|
class Vocos(nn.Module): |
|
def __init__( |
|
self, |
|
input_channels: int = 128, |
|
dim: int = 512, |
|
intermediate_dim: int = 4096, |
|
num_layers: int = 30, |
|
n_fft: int = 640, |
|
hop_size: int = 160, |
|
padding: str = "same", |
|
adanorm_num_embeddings=None, |
|
): |
|
super().__init__() |
|
|
|
self.backbone = VocosBackbone( |
|
input_channels=input_channels, |
|
dim=dim, |
|
intermediate_dim=intermediate_dim, |
|
num_layers=num_layers, |
|
adanorm_num_embeddings=adanorm_num_embeddings, |
|
) |
|
self.head = ISTFTHead(dim, n_fft, hop_size, padding) |
|
self.hop_size = hop_size |
|
|
|
def forward(self, x, input_length): |
|
x = self.backbone(x) |
|
x = self.head(x) |
|
output_length = input_length * self.hop_size |
|
return x[:, None, :], output_length |
|
|
|
|