421 lines
15 KiB
Python
421 lines
15 KiB
Python
# Extracted from transformers' WhisperModel to simplify package dependency
|
|
import math
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from typing import Optional, Literal
|
|
from fireredtts2.codec.utils import make_nonpad_mask
|
|
from fireredtts2.codec.audio import mel_filter_bank
|
|
|
|
|
|
def sinusoids(length: int, channels: int, max_timescale: float = 10000) -> torch.Tensor:
|
|
"""Returns sinusoids for positional embedding"""
|
|
if channels % 2 != 0:
|
|
raise ValueError(
|
|
f"Number of channels has to be divisible by 2 for sinusoidal positional embeddings, got {channels} channels."
|
|
)
|
|
log_timescale_increment = math.log(max_timescale) / (channels // 2 - 1)
|
|
inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))
|
|
scaled_time = torch.arange(length).view(-1, 1) * inv_timescales.view(1, -1)
|
|
return torch.cat([scaled_time.sin(), scaled_time.cos()], dim=1)
|
|
|
|
|
|
class WhisperSdpaAttention(nn.Module):
|
|
def __init__(
|
|
self,
|
|
embed_dim: int,
|
|
num_heads: int,
|
|
dropout: float = 0.0,
|
|
bias: bool = True,
|
|
):
|
|
super().__init__()
|
|
self.embed_dim = embed_dim
|
|
self.num_heads = num_heads
|
|
self.dropout = dropout
|
|
self.head_dim = embed_dim // num_heads
|
|
self.bias = bias
|
|
self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
|
|
self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
|
self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
|
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
|
|
|
|
def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
|
|
return (
|
|
tensor.view(bsz, seq_len, self.num_heads, self.head_dim)
|
|
.transpose(1, 2)
|
|
.contiguous()
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
):
|
|
"""
|
|
Args:
|
|
attention_mask: Bool mask or float mask. Bool mask, True indicates should attend. Float mask is added to the attention score.
|
|
"""
|
|
bsz, tgt_len, _ = hidden_states.size()
|
|
|
|
query_states = self._shape(self.q_proj(hidden_states), tgt_len, bsz)
|
|
key_states = self._shape(self.k_proj(hidden_states), tgt_len, bsz)
|
|
value_states = self._shape(self.v_proj(hidden_states), tgt_len, bsz)
|
|
|
|
# NOTE sdpa needs a 4-dim attention_mask: (b, nh, tq, tv)
|
|
if attention_mask is not None and len(attention_mask.shape) == 3:
|
|
attention_mask = attention_mask.unsqueeze(1)
|
|
|
|
attn_output = F.scaled_dot_product_attention(
|
|
query_states,
|
|
key_states,
|
|
value_states,
|
|
attn_mask=attention_mask,
|
|
dropout_p=self.dropout if self.training else 0.0,
|
|
) # (bsz, nh, l, d)
|
|
attn_output = attn_output.transpose(1, 2)
|
|
attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
|
|
|
|
attn_output = self.out_proj(attn_output)
|
|
return attn_output
|
|
|
|
def forward_chunk(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
kv_cache: torch.Tensor = None,
|
|
):
|
|
"""Forward self-attention with kv cache.
|
|
|
|
Args:
|
|
hidden_states: shape (b, t, c)
|
|
kv_cache: shape (b, nh, t, c*2)
|
|
"""
|
|
bsz, tgt_len, _ = hidden_states.size()
|
|
|
|
# shape (b, nh, t, c)
|
|
query_states = self._shape(self.q_proj(hidden_states), tgt_len, bsz)
|
|
key_states = self._shape(self.k_proj(hidden_states), tgt_len, bsz)
|
|
value_states = self._shape(self.v_proj(hidden_states), tgt_len, bsz)
|
|
|
|
# unpack cache
|
|
if kv_cache is not None:
|
|
k_cache, v_cache = kv_cache.chunk(2, dim=-1)
|
|
key_states = torch.cat([k_cache, key_states], dim=2)
|
|
value_states = torch.cat([v_cache, value_states], dim=2)
|
|
new_kv_cache = torch.cat([key_states, value_states], dim=-1)
|
|
|
|
# attention
|
|
attn_output = F.scaled_dot_product_attention(
|
|
query_states,
|
|
key_states,
|
|
value_states,
|
|
attn_mask=None,
|
|
dropout_p=0.0,
|
|
) # (bsz, nh, l, d)
|
|
attn_output = attn_output.transpose(1, 2)
|
|
attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
|
|
|
|
attn_output = self.out_proj(attn_output)
|
|
return attn_output, new_kv_cache
|
|
|
|
|
|
class WhisperEncoderLayer(nn.Module):
|
|
def __init__(
|
|
self,
|
|
embed_dim: int,
|
|
num_heads: int,
|
|
ffn_dim: int = None,
|
|
attn_dropout: float = 0.0,
|
|
dropout: float = 0.0,
|
|
):
|
|
super().__init__()
|
|
self.dropout = dropout
|
|
# Attention
|
|
self.self_attn = WhisperSdpaAttention(embed_dim, num_heads, attn_dropout)
|
|
self.self_attn_layer_norm = nn.LayerNorm(embed_dim)
|
|
# FFN
|
|
ffn_dim = ffn_dim if ffn_dim is not None else embed_dim * 4
|
|
self.fc1 = nn.Linear(embed_dim, ffn_dim)
|
|
self.fc2 = nn.Linear(ffn_dim, embed_dim)
|
|
# Output norm
|
|
self.final_layer_norm = nn.LayerNorm(embed_dim)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
attention_mask: torch.Tensor,
|
|
):
|
|
# Attention
|
|
residual = hidden_states
|
|
hidden_states = self.self_attn_layer_norm(hidden_states)
|
|
hidden_states = self.self_attn(hidden_states, attention_mask)
|
|
hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
|
|
hidden_states = residual + hidden_states
|
|
|
|
# FFN
|
|
residual = hidden_states
|
|
hidden_states = self.final_layer_norm(hidden_states)
|
|
hidden_states = F.gelu(self.fc1(hidden_states))
|
|
hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
|
|
hidden_states = self.fc2(hidden_states)
|
|
hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
|
|
hidden_states = residual + hidden_states
|
|
return hidden_states
|
|
|
|
def forward_chunk(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
kv_cache: torch.Tensor = None,
|
|
):
|
|
"""Forward self-attention with kv cache.
|
|
|
|
Args:
|
|
hidden_states: shape (b, t, c)
|
|
kv_cache: shape (b, nh, t, c*2)
|
|
"""
|
|
# Attention
|
|
residual = hidden_states
|
|
hidden_states = self.self_attn_layer_norm(hidden_states)
|
|
hidden_states, new_kv_cache = self.self_attn.forward_chunk(
|
|
hidden_states, kv_cache
|
|
)
|
|
hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
|
|
hidden_states = residual + hidden_states
|
|
|
|
# FFN
|
|
residual = hidden_states
|
|
hidden_states = self.final_layer_norm(hidden_states)
|
|
hidden_states = F.gelu(self.fc1(hidden_states))
|
|
hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
|
|
hidden_states = self.fc2(hidden_states)
|
|
hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
|
|
hidden_states = residual + hidden_states
|
|
return hidden_states, new_kv_cache
|
|
|
|
|
|
class WhisperEncoder(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_dim: int,
|
|
embed_dim: int,
|
|
num_layers: int,
|
|
num_heads: int,
|
|
ffn_dim: int = None,
|
|
attn_dropout: float = 0.0,
|
|
dropout: float = 0.0,
|
|
max_positions: int = 1500,
|
|
):
|
|
super().__init__()
|
|
self.in_dim = in_dim
|
|
self.embed_dim = embed_dim
|
|
self.dropout = dropout
|
|
# Input downsampling
|
|
self.conv1 = nn.Conv1d(in_dim, embed_dim, kernel_size=3, padding=1)
|
|
self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1)
|
|
# Fixed positional embedding
|
|
self.max_positions = max_positions
|
|
self.embed_positions = nn.Embedding(self.max_positions, embed_dim)
|
|
self.embed_positions.requires_grad_(False)
|
|
# Transformer
|
|
self.layers = nn.ModuleList(
|
|
[
|
|
WhisperEncoderLayer(
|
|
embed_dim, num_heads, ffn_dim, attn_dropout, dropout
|
|
)
|
|
for _ in range(num_layers)
|
|
]
|
|
)
|
|
# Output norm
|
|
self.layer_norm = nn.LayerNorm(embed_dim)
|
|
# Init weight
|
|
self.apply(self._init_weights)
|
|
# Init position embedding
|
|
self.embed_positions.weight.copy_(sinusoids(*self.embed_positions.weight.shape))
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
hidden_length: torch.Tensor,
|
|
apply_position: bool = True,
|
|
):
|
|
# Downsampling
|
|
hidden_states = hidden_states.transpose(1, 2)
|
|
hidden_states = F.gelu(self.conv1(hidden_states))
|
|
hidden_states = F.gelu(self.conv2(hidden_states))
|
|
hidden_states = hidden_states.transpose(1, 2)
|
|
hidden_length = hidden_length // 2 # from 100Hz -> 50Hz
|
|
# Pos encoding
|
|
if apply_position:
|
|
pos_embed = self.embed_positions(
|
|
torch.arange(0, hidden_states.shape[1], device=hidden_states.device)
|
|
)
|
|
hidden_states = hidden_states + pos_embed
|
|
hidden_states = nn.functional.dropout(
|
|
hidden_states, p=self.dropout, training=self.training
|
|
)
|
|
# Transformer
|
|
attention_mask = make_nonpad_mask(hidden_length).unsqueeze(1) # (b, 1, t)
|
|
for layer in self.layers:
|
|
hidden_states = layer(hidden_states, attention_mask)
|
|
|
|
hidden_states = self.layer_norm(hidden_states)
|
|
return hidden_states, hidden_length
|
|
|
|
def _init_weights(self, module):
|
|
std = 0.02
|
|
if isinstance(module, (nn.Linear, nn.Conv1d)):
|
|
module.weight.data.normal_(mean=0.0, std=std)
|
|
if module.bias is not None:
|
|
module.bias.data.zero_()
|
|
elif isinstance(module, nn.Embedding):
|
|
module.weight.data.normal_(mean=0.0, std=std)
|
|
if module.padding_idx is not None:
|
|
module.weight.data[module.padding_idx].zero_()
|
|
|
|
|
|
class WhisperMelExtractor(nn.Module):
|
|
def __init__(
|
|
self,
|
|
num_mels: int = 128,
|
|
sampling_rate: int = 16000,
|
|
hop_length: int = 160,
|
|
n_fft: int = 400,
|
|
fmin: float = 0,
|
|
fmax: float = 8000,
|
|
padding_value=0.0,
|
|
):
|
|
super().__init__()
|
|
self.num_mels = num_mels
|
|
self.sampling_rate = sampling_rate
|
|
self.hop_length = hop_length
|
|
self.n_fft = n_fft
|
|
self.fmin = fmin
|
|
self.fmax = fmax
|
|
self.padding_value = padding_value
|
|
self.mel_filters = mel_filter_bank(
|
|
num_frequency_bins=(1 + n_fft // 2),
|
|
num_mel_filters=num_mels,
|
|
min_frequency=fmin,
|
|
max_frequency=fmax,
|
|
sampling_rate=sampling_rate,
|
|
norm="slaney",
|
|
mel_scale="slaney",
|
|
)
|
|
|
|
def extract_fbank(self, audio: torch.Tensor):
|
|
"""
|
|
Args:
|
|
audio: batched audio of shape (b, t)
|
|
"""
|
|
device = audio.device # compute on cuda if input is on cuda
|
|
# Mel
|
|
window = torch.hann_window(self.n_fft).to(device)
|
|
stft = torch.stft(
|
|
audio, self.n_fft, self.hop_length, window=window, return_complex=True
|
|
)
|
|
magnitudes = stft[..., :-1].abs() ** 2
|
|
mel_filters = torch.from_numpy(self.mel_filters).type(torch.float32).to(device)
|
|
mel_spec = mel_filters.T @ magnitudes
|
|
log_spec = torch.clamp(mel_spec, min=1e-10).log10()
|
|
# Norm
|
|
max_val = log_spec.max(dim=2, keepdim=True)[0].max(dim=1, keepdim=True)[0]
|
|
log_spec = torch.maximum(log_spec, max_val - 8.0)
|
|
log_spec = (log_spec + 4.0) / 4.0
|
|
return log_spec
|
|
|
|
def __call__(self, audio16k: torch.Tensor, audio16k_length: torch.Tensor):
|
|
mel = self.extract_fbank(audio16k).transpose(1, 2)
|
|
mel_length = audio16k_length // self.hop_length
|
|
# mel: (b, t, c=128)
|
|
return mel, mel_length
|
|
|
|
|
|
# Pretrained encoder from whisper-large-v3
|
|
class PretrainedWhisperEncoder(WhisperEncoder):
|
|
@classmethod
|
|
def from_pretrained(cls, pretrained_path: str = None):
|
|
encoder = cls(
|
|
in_dim=128,
|
|
embed_dim=1280,
|
|
num_layers=32,
|
|
num_heads=20,
|
|
ffn_dim=5120,
|
|
attn_dropout=0.0,
|
|
max_positions=1500,
|
|
)
|
|
if pretrained_path is not None:
|
|
ckpt = torch.load(pretrained_path, map_location="cpu")
|
|
encoder.load_state_dict(ckpt)
|
|
encoder.eval()
|
|
# Disable grad
|
|
for p in encoder.parameters():
|
|
p.requires_grad_(False)
|
|
# Add Mel extractor
|
|
encoder.feature_extractor = WhisperMelExtractor(
|
|
num_mels=128,
|
|
sampling_rate=16000,
|
|
hop_length=160,
|
|
n_fft=400,
|
|
fmin=0,
|
|
fmax=8000,
|
|
)
|
|
return encoder
|
|
|
|
@torch.inference_mode()
|
|
def forward(self, audio16k: torch.Tensor, audio16k_length: torch.Tensor):
|
|
# Extract mel
|
|
mel, mel_length = self.feature_extractor(audio16k, audio16k_length)
|
|
# Forward model
|
|
semantic_feats, semantic_length = super().forward(
|
|
mel, mel_length, apply_position=True
|
|
)
|
|
return semantic_feats, semantic_length
|
|
|
|
|
|
class WhisperAcousticEncoder(WhisperEncoder):
|
|
def __init__(
|
|
self,
|
|
# Mel extraction params
|
|
num_mels: int = 128,
|
|
sampling_rate: int = 16000,
|
|
hop_length: int = 160,
|
|
n_fft: int = 400,
|
|
fmin: float = 0.0,
|
|
fmax: float = 8000,
|
|
# Encoder params
|
|
embed_dim: int = 768,
|
|
num_layers: int = 12,
|
|
num_heads: int = 8,
|
|
ffn_dim: int = None,
|
|
attn_dropout: float = 0.0,
|
|
dropout: float = 0.0,
|
|
max_positions: int = 1500, # 50Hz * 30s
|
|
):
|
|
super().__init__(
|
|
in_dim=num_mels,
|
|
embed_dim=embed_dim,
|
|
num_layers=num_layers,
|
|
num_heads=num_heads,
|
|
ffn_dim=ffn_dim,
|
|
attn_dropout=attn_dropout,
|
|
dropout=dropout,
|
|
max_positions=max_positions,
|
|
)
|
|
self.feature_extractor = WhisperMelExtractor(
|
|
num_mels=num_mels,
|
|
sampling_rate=sampling_rate,
|
|
hop_length=hop_length,
|
|
n_fft=n_fft,
|
|
fmin=fmin,
|
|
fmax=fmax,
|
|
)
|
|
|
|
def forward(self, audio16k: torch.Tensor, audio16k_length: torch.Tensor):
|
|
# Extract mel
|
|
with torch.no_grad():
|
|
mel, mel_length = self.feature_extractor(audio16k, audio16k_length)
|
|
# Forward model
|
|
hidden_states, hidden_length = super().forward(
|
|
mel, mel_length, apply_position=True
|
|
)
|
|
return hidden_states, hidden_length
|