229 lines
8.3 KiB
Python
229 lines
8.3 KiB
Python
|
# -*- encoding: utf-8 -*-
|
||
|
|
# File: audio.py
|
||
|
|
# Description: None
|
||
|
|
|
||
|
|
|
||
|
|
from typing import Iterable, List, Optional
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
import torch.nn.functional as F
|
||
|
|
from torch import Tensor
|
||
|
|
|
||
|
|
|
||
|
|
class LayerNorm(nn.LayerNorm):
|
||
|
|
def forward(self, x: Tensor) -> Tensor:
|
||
|
|
return super().forward(x).type(x.dtype)
|
||
|
|
|
||
|
|
|
||
|
|
class Linear(nn.Linear):
|
||
|
|
def forward(self, x: Tensor) -> Tensor:
|
||
|
|
return F.linear(
|
||
|
|
x,
|
||
|
|
self.weight.to(x.dtype),
|
||
|
|
None if self.bias is None else self.bias.to(x.dtype),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class Conv1d(nn.Conv1d):
|
||
|
|
def _conv_forward(self, x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> Tensor:
|
||
|
|
return super()._conv_forward(x, weight.to(x.dtype), None if bias is None else bias.to(x.dtype))
|
||
|
|
|
||
|
|
|
||
|
|
def sinusoids(length, channels, max_timescale=10000):
|
||
|
|
"""Returns sinusoids 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)
|
||
|
|
|
||
|
|
|
||
|
|
class MultiHeadAttention(nn.Module):
|
||
|
|
def __init__(self, n_state: int, n_head: int):
|
||
|
|
super().__init__()
|
||
|
|
self.n_head = n_head
|
||
|
|
self.query = Linear(n_state, n_state)
|
||
|
|
self.key = Linear(n_state, n_state, bias=False)
|
||
|
|
self.value = Linear(n_state, n_state)
|
||
|
|
self.out = Linear(n_state, n_state)
|
||
|
|
|
||
|
|
def forward(
|
||
|
|
self,
|
||
|
|
x: Tensor,
|
||
|
|
xa: Optional[Tensor] = None,
|
||
|
|
mask: Optional[Tensor] = None,
|
||
|
|
kv_cache: Optional[dict] = None,
|
||
|
|
):
|
||
|
|
q = self.query(x)
|
||
|
|
|
||
|
|
if kv_cache is None or xa is None or self.key not in kv_cache:
|
||
|
|
# hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors;
|
||
|
|
# otherwise, perform key/value projections for self- or cross-attention as usual.
|
||
|
|
k = self.key(x if xa is None else xa)
|
||
|
|
v = self.value(x if xa is None else xa)
|
||
|
|
else:
|
||
|
|
# for cross-attention, calculate keys and values once and reuse in subsequent calls.
|
||
|
|
k = kv_cache[self.key]
|
||
|
|
v = kv_cache[self.value]
|
||
|
|
|
||
|
|
wv, qk = self.qkv_attention(q, k, v, mask)
|
||
|
|
return self.out(wv), qk
|
||
|
|
|
||
|
|
def qkv_attention(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None):
|
||
|
|
n_batch, n_ctx, n_state = q.shape
|
||
|
|
scale = (n_state // self.n_head) ** -0.25
|
||
|
|
q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) * scale
|
||
|
|
k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 3, 1) * scale
|
||
|
|
v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
|
||
|
|
|
||
|
|
qk = q @ k
|
||
|
|
if mask is not None:
|
||
|
|
qk += mask
|
||
|
|
|
||
|
|
w = F.softmax(qk, dim=-1).to(q.dtype)
|
||
|
|
return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2), qk.detach()
|
||
|
|
|
||
|
|
|
||
|
|
class ResidualAttentionBlock(nn.Module):
|
||
|
|
def __init__(self, n_state: int, n_head: int, cross_attention: bool = False):
|
||
|
|
super().__init__()
|
||
|
|
|
||
|
|
self.attn = MultiHeadAttention(n_state, n_head)
|
||
|
|
self.attn_ln = LayerNorm(n_state)
|
||
|
|
|
||
|
|
self.cross_attn = MultiHeadAttention(n_state, n_head) if cross_attention else None
|
||
|
|
self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None
|
||
|
|
|
||
|
|
n_mlp = n_state * 4
|
||
|
|
self.mlp = nn.Sequential(Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state))
|
||
|
|
self.mlp_ln = LayerNorm(n_state)
|
||
|
|
|
||
|
|
def forward(
|
||
|
|
self,
|
||
|
|
x: Tensor,
|
||
|
|
xa: Optional[Tensor] = None,
|
||
|
|
mask: Optional[Tensor] = None,
|
||
|
|
kv_cache: Optional[dict] = None,
|
||
|
|
):
|
||
|
|
x = x + self.attn(self.attn_ln(x), mask=mask, kv_cache=kv_cache)[0]
|
||
|
|
if self.cross_attn:
|
||
|
|
x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache)[0]
|
||
|
|
x = x + self.mlp(self.mlp_ln(x))
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
class AudioEncoder(nn.Module):
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
n_mels: int,
|
||
|
|
n_ctx: int,
|
||
|
|
n_state: int,
|
||
|
|
n_head: int,
|
||
|
|
n_layer: int,
|
||
|
|
output_dim: int = 512,
|
||
|
|
avg_pool: bool = True,
|
||
|
|
add_audio_bos_eos_token: bool = True,
|
||
|
|
**kwargs,
|
||
|
|
):
|
||
|
|
super().__init__()
|
||
|
|
self.conv1 = Conv1d(n_mels, n_state, kernel_size=3, padding=1)
|
||
|
|
self.conv2 = Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1)
|
||
|
|
self.register_buffer("positional_embedding", sinusoids(n_ctx, n_state))
|
||
|
|
|
||
|
|
self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList(
|
||
|
|
[ResidualAttentionBlock(n_state, n_head) for _ in range(n_layer)]
|
||
|
|
)
|
||
|
|
self.ln_post = LayerNorm(n_state)
|
||
|
|
|
||
|
|
if avg_pool:
|
||
|
|
self.avg_pooler = nn.AvgPool1d(2, stride=2)
|
||
|
|
else:
|
||
|
|
self.avg_pooler = None
|
||
|
|
self.proj = nn.Linear(n_state, output_dim)
|
||
|
|
if add_audio_bos_eos_token:
|
||
|
|
self.audio_bos_eos_token = nn.Embedding(2, output_dim)
|
||
|
|
else:
|
||
|
|
self.audio_bos_eos_token = None
|
||
|
|
self.output_dim = output_dim
|
||
|
|
self.n_head = n_head
|
||
|
|
|
||
|
|
def forward(self, x: Tensor, padding_mask: Tensor = None, audio_lengths: Tensor = None):
|
||
|
|
"""
|
||
|
|
x : torch.Tensor, shape = (batch_size, n_mels, n_ctx)
|
||
|
|
the mel spectrogram of the audio
|
||
|
|
"""
|
||
|
|
x = x.to(dtype=self.conv1.weight.dtype, device=self.conv1.weight.device)
|
||
|
|
if audio_lengths is not None:
|
||
|
|
input_mel_len = audio_lengths[:, 0] * 2
|
||
|
|
max_mel_len_in_batch = input_mel_len.max()
|
||
|
|
x = x[:, :, :max_mel_len_in_batch]
|
||
|
|
x = F.gelu(self.conv1(x))
|
||
|
|
x = F.gelu(self.conv2(x))
|
||
|
|
x = x.permute(0, 2, 1) # B, L, D
|
||
|
|
bsz = x.size(0)
|
||
|
|
src_len = x.size(1)
|
||
|
|
|
||
|
|
self.input_positional_embedding = self.positional_embedding[:src_len]
|
||
|
|
assert (
|
||
|
|
x.shape[1:] == self.input_positional_embedding.shape
|
||
|
|
), f"incorrect audio shape: {x.shape[1:], self.input_positional_embedding.shape}"
|
||
|
|
x = (x + self.input_positional_embedding).to(x.dtype)
|
||
|
|
if padding_mask is not None:
|
||
|
|
padding_mask = padding_mask.to(dtype=self.conv1.weight.dtype, device=self.conv1.weight.device)
|
||
|
|
batch_src_len = padding_mask.size(1)
|
||
|
|
x = x[:, :batch_src_len, :]
|
||
|
|
padding_mask = padding_mask.view(bsz, -1, batch_src_len)
|
||
|
|
padding_mask_ = padding_mask.all(1)
|
||
|
|
x[padding_mask_] = 0
|
||
|
|
key_padding_mask = (
|
||
|
|
padding_mask_.view(bsz, 1, 1, batch_src_len)
|
||
|
|
.expand(-1, self.n_head, -1, -1)
|
||
|
|
.reshape(bsz, self.n_head, 1, batch_src_len)
|
||
|
|
)
|
||
|
|
new_padding_mask = torch.zeros_like(key_padding_mask, dtype=x.dtype)
|
||
|
|
padding_mask = new_padding_mask.masked_fill(key_padding_mask, float("-inf"))
|
||
|
|
|
||
|
|
for block in self.blocks:
|
||
|
|
x = block(x, mask=padding_mask)
|
||
|
|
|
||
|
|
if self.avg_pooler:
|
||
|
|
x = x.permute(0, 2, 1)
|
||
|
|
x = self.avg_pooler(x)
|
||
|
|
x = x.permute(0, 2, 1)
|
||
|
|
|
||
|
|
x = self.ln_post(x)
|
||
|
|
x = self.proj(x)
|
||
|
|
|
||
|
|
if self.audio_bos_eos_token is not None:
|
||
|
|
bos = self.audio_bos_eos_token.weight[0][None, :]
|
||
|
|
eos = self.audio_bos_eos_token.weight[1][None, :]
|
||
|
|
else:
|
||
|
|
bos, eos = None, None
|
||
|
|
return x, bos, eos
|
||
|
|
|
||
|
|
def encode(
|
||
|
|
self,
|
||
|
|
input_audios: Tensor,
|
||
|
|
input_audio_lengths: Tensor,
|
||
|
|
audio_span_tokens: List,
|
||
|
|
):
|
||
|
|
real_input_audio_lens = input_audio_lengths[:, 0].tolist()
|
||
|
|
max_len_in_batch = max(real_input_audio_lens)
|
||
|
|
padding_mask = torch.ones([input_audios.size(0), max_len_in_batch]).to(
|
||
|
|
dtype=self.conv1.weight.dtype, device=self.conv1.weight.device
|
||
|
|
)
|
||
|
|
for index in range(len(input_audios)):
|
||
|
|
padding_mask[index, : input_audio_lengths[index][0].item()] = 0
|
||
|
|
x, bos, eos = self(input_audios, padding_mask, input_audio_lengths)
|
||
|
|
output_audios = []
|
||
|
|
for i in range(len(audio_span_tokens)):
|
||
|
|
audio_span = audio_span_tokens[i]
|
||
|
|
audio = x[i][: audio_span - 2]
|
||
|
|
if bos is not None:
|
||
|
|
audio = torch.concat([bos, audio, eos])
|
||
|
|
assert len(audio) == audio_span
|
||
|
|
output_audios.append(audio)
|
||
|
|
return output_audios
|