377 lines
13 KiB
Python
377 lines
13 KiB
Python
import math
|
|
import json
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from typing import List, Dict
|
|
from torch.nn.utils.rnn import pad_sequence
|
|
|
|
from fireredtts2.codec.rvq import ResidualVQ
|
|
from fireredtts2.codec.decoder import AcousticDecoder
|
|
from fireredtts2.codec.utils import make_nonpad_mask
|
|
from fireredtts2.codec.whisper import (
|
|
WhisperEncoderLayer,
|
|
PretrainedWhisperEncoder,
|
|
WhisperAcousticEncoder,
|
|
)
|
|
|
|
|
|
class SslAdaptor(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_dim: int,
|
|
embed_dim: int,
|
|
out_dim: int,
|
|
num_layers: int,
|
|
num_heads: int,
|
|
ffn_dim: int = None,
|
|
attn_dropout: float = 0.0,
|
|
dropout: float = 0.0,
|
|
):
|
|
super().__init__()
|
|
self.in_dim = in_dim
|
|
self.embed_dim = embed_dim
|
|
self.dropout = dropout
|
|
# Input Projection
|
|
self.in_proj = nn.Linear(in_dim, embed_dim)
|
|
# 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)
|
|
# Output projection
|
|
self.out_proj = nn.Linear(embed_dim, out_dim)
|
|
# Init weight
|
|
self.apply(self._init_weights)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
hidden_length: torch.Tensor,
|
|
):
|
|
# Downsampling
|
|
hidden_states = self.in_proj(hidden_states)
|
|
# 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)
|
|
hidden_states = self.out_proj(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 ResidualDownConv(nn.Module):
|
|
def __init__(
|
|
self,
|
|
embed_dim: int = 768,
|
|
avg_pooler=4,
|
|
):
|
|
super().__init__()
|
|
self.embed_dim = embed_dim
|
|
self.avg_pooler = avg_pooler
|
|
self.intermediate_dim = embed_dim * avg_pooler
|
|
# Convolution layer for downsampling
|
|
self.gate_proj = nn.Conv1d(
|
|
embed_dim, self.intermediate_dim, avg_pooler, avg_pooler, bias=False
|
|
)
|
|
self.up_proj = nn.Conv1d(
|
|
embed_dim, self.intermediate_dim, avg_pooler, avg_pooler, bias=False
|
|
)
|
|
# Downsampled linear projection
|
|
self.down_proj = nn.Linear(
|
|
self.intermediate_dim, self.intermediate_dim, bias=False
|
|
)
|
|
# Activation function and layer normalization
|
|
self.act_fn = nn.SiLU()
|
|
self.layer_norm = nn.LayerNorm(self.intermediate_dim)
|
|
# Final output projection
|
|
self.out_proj = nn.Linear(self.intermediate_dim, embed_dim)
|
|
|
|
def forward(self, x: torch.Tensor, input_length: torch.Tensor):
|
|
output_length = input_length // self.avg_pooler
|
|
batch_size, seq_len, _ = x.shape # (B, T, D)
|
|
|
|
xt = x.permute(0, 2, 1) # (B, D, T)
|
|
g = self.gate_proj(xt).permute(0, 2, 1) # (B, T//4, D*4)
|
|
u = self.up_proj(xt).permute(0, 2, 1) # (B, T//4, D*4)
|
|
x = x.reshape(batch_size, -1, self.intermediate_dim) # (B, T//4, D*4)
|
|
|
|
c = self.down_proj(self.act_fn(g) * u) # (B, T//4, D*4)
|
|
res = self.layer_norm(c + x) # (B, T//4, D*4)
|
|
|
|
res = self.out_proj(res)
|
|
return res, output_length
|
|
|
|
|
|
class UpConv(nn.Module):
|
|
def __init__(
|
|
self,
|
|
embed_dim: int = 768,
|
|
stride: int = 4,
|
|
):
|
|
super().__init__()
|
|
self.embed_dim = embed_dim
|
|
self.stride = stride
|
|
self.in_proj = nn.Linear(embed_dim, self.stride * embed_dim)
|
|
# Simple transpose convolution layer to keep channel number consistent
|
|
self.up_conv = nn.ConvTranspose1d(
|
|
self.stride * embed_dim,
|
|
embed_dim,
|
|
kernel_size=stride,
|
|
stride=stride,
|
|
bias=False,
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor, input_length: torch.Tensor):
|
|
x = self.in_proj(x)
|
|
x = x.transpose(1, 2)
|
|
res = self.up_conv(x)
|
|
res = res.transpose(1, 2)
|
|
output_length = input_length * self.stride
|
|
return res, output_length
|
|
|
|
|
|
class RedCodec(nn.Module):
|
|
def __init__(
|
|
self,
|
|
ssl: PretrainedWhisperEncoder,
|
|
ssl_adaptor: SslAdaptor,
|
|
acoustic_encoder: WhisperAcousticEncoder,
|
|
downsample: ResidualDownConv,
|
|
rvq: ResidualVQ,
|
|
upsample: UpConv,
|
|
semantic_decoder: SslAdaptor,
|
|
acoustic_decoder: AcousticDecoder,
|
|
):
|
|
super().__init__()
|
|
self.ssl = ssl
|
|
self.ssl_adaptor = ssl_adaptor
|
|
self.acoustic_encoder = acoustic_encoder
|
|
self.downsample = downsample
|
|
self.rvq = rvq
|
|
self.upsample = upsample
|
|
self.semantic_decoder = semantic_decoder
|
|
self.acoustic_decoder = acoustic_decoder
|
|
|
|
@classmethod
|
|
def from_config(cls, config_json: str) -> "RedCodec":
|
|
with open(config_json, "rb") as f:
|
|
config = json.load(f)["codec"]
|
|
ssl = PretrainedWhisperEncoder.from_pretrained()
|
|
ssl_adaptor = SslAdaptor(**config["ssl_adaptor"])
|
|
acoustic_encoder = WhisperAcousticEncoder(**config["acoustic_encoder"])
|
|
downsample = ResidualDownConv(**config["downsample"])
|
|
rvq = ResidualVQ(**config["rvq"])
|
|
upsample = UpConv(**config["upsample"])
|
|
semantic_decoder = SslAdaptor(**config["semantic_decoder"])
|
|
acoustic_decoder = AcousticDecoder(**config["acoustic_decoder"])
|
|
return cls(
|
|
ssl,
|
|
ssl_adaptor,
|
|
acoustic_encoder,
|
|
downsample,
|
|
rvq,
|
|
upsample,
|
|
semantic_decoder,
|
|
acoustic_decoder,
|
|
)
|
|
|
|
|
|
class RedCodecInfer(RedCodec):
|
|
def __init__(self, codec: RedCodec):
|
|
super().__init__(
|
|
codec.ssl,
|
|
codec.ssl_adaptor,
|
|
codec.acoustic_encoder,
|
|
codec.downsample,
|
|
codec.rvq,
|
|
codec.upsample,
|
|
codec.semantic_decoder,
|
|
codec.acoustic_decoder,
|
|
)
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, conf_path: str, ckpt_path: str) -> "RedCodecInfer":
|
|
with open(conf_path, "r") as f:
|
|
codec = RedCodec.from_config(conf_path)
|
|
ckpt = torch.load(ckpt_path)["generator"]
|
|
codec.load_state_dict(ckpt)
|
|
return cls(codec)
|
|
|
|
def _encode_one_batch(self, audio16k: torch.Tensor):
|
|
B, T = audio16k.shape
|
|
audio16k_length = torch.tensor(
|
|
[T] * B, dtype=torch.long, device=audio16k.device
|
|
)
|
|
# Semantic
|
|
ssl, ssl_length = self.ssl.forward(audio16k, audio16k_length)
|
|
ssl = ssl.clone() # For onnx export
|
|
sem_feats, sem_length = self.ssl_adaptor(ssl, ssl_length)
|
|
# Acoustic
|
|
aco_feats, aco_length = self.acoustic_encoder(audio16k, audio16k_length)
|
|
# VQ
|
|
vq_in_feats = torch.cat([sem_feats, aco_feats], dim=2)
|
|
vq_in_feats, vq_in_length = self.downsample(vq_in_feats, aco_length)
|
|
# RVQ,
|
|
indices = self.rvq.encode_codes(vq_in_feats.transpose(1, 2)) # (nq, B, L)
|
|
indices = indices.permute(1, 0, 2)
|
|
return indices # (B, nq, L)
|
|
|
|
@staticmethod
|
|
def _pad_and_chunk(audio: torch.Tensor, chunk_size: int) -> List[torch.Tensor]:
|
|
pad_len = math.ceil(audio.shape[1] / chunk_size) * chunk_size - audio.shape[1]
|
|
audio = F.pad(audio, (0, pad_len), mode="constant", value=0)
|
|
audio_chunks = audio.split(chunk_size, dim=1)
|
|
return audio_chunks
|
|
|
|
@torch.inference_mode()
|
|
def encode(
|
|
self,
|
|
audio16k: torch.Tensor,
|
|
audio16k_length: torch.Tensor = None,
|
|
batch_size: int = 96,
|
|
):
|
|
"""
|
|
Args:
|
|
audio16k: shape (b, t)
|
|
audio16k_length: (b,)
|
|
Returns:
|
|
token: shape (b, nq, l)
|
|
token_length: (b,)
|
|
"""
|
|
if audio16k_length is None:
|
|
assert audio16k.shape[0] == 1
|
|
audio16k_length = torch.tensor(
|
|
[audio16k.shape[1]], dtype=torch.long, device=audio16k.device
|
|
)
|
|
|
|
CHUNK_SIZE = 6 * 16000
|
|
B, T = audio16k.shape
|
|
# Pad, chunk, and batch
|
|
audio16k_batch = []
|
|
batch_size_list = []
|
|
for i in range(B):
|
|
# Remove extra paddings
|
|
one_audio_chunks = self._pad_and_chunk(
|
|
audio16k[i : (i + 1), : audio16k_length[i]], CHUNK_SIZE
|
|
)
|
|
audio16k_batch += one_audio_chunks
|
|
batch_size_list.append(len(one_audio_chunks))
|
|
audio16k_batch = torch.cat(audio16k_batch, dim=0)
|
|
# Batch encode
|
|
token_batch = []
|
|
for i in range(0, audio16k_batch.shape[0], batch_size):
|
|
one_audio_batch = audio16k_batch[i : (i + batch_size)]
|
|
one_token_batch = self._encode_one_batch(one_audio_batch)
|
|
token_batch.append(one_token_batch)
|
|
token_batch = torch.cat(token_batch, dim=0)
|
|
# Recover & concat
|
|
token_list = torch.split(
|
|
token_batch, batch_size_list, dim=0
|
|
) # [(B=1, nq, l), (B=3, nq, l), ...]
|
|
token_list = [
|
|
torch.cat(token_ts.split(1, dim=0), dim=-1) # (B=1, nq, l)
|
|
for token_ts in token_list
|
|
]
|
|
# Pad tokens
|
|
token = pad_sequence(
|
|
[ts.squeeze(0).transpose(1, 0) for ts in token_list],
|
|
batch_first=True,
|
|
padding_value=0,
|
|
).transpose(
|
|
1, 2
|
|
) # (B, nq, L)
|
|
token_length = (audio16k_length / 1280).ceil().long()
|
|
token = token[
|
|
..., : token_length.max()
|
|
] # Remove extra paddings (we pad to multiples of 6s)
|
|
return token, token_length
|
|
|
|
@torch.inference_mode()
|
|
def decode(self, tokens: torch.Tensor):
|
|
"""
|
|
Args:
|
|
tokens: (B=1, nq, L)
|
|
Returns:
|
|
audio: (B=1, t)
|
|
"""
|
|
tokens = tokens.permute(1, 0, 2) # (B, nq, L) -> (nq, B, L)
|
|
vq_out_feats = self.rvq.decode_codes(tokens)
|
|
vq_out_feats = vq_out_feats.transpose(1, 2)
|
|
vq_out_length = torch.tensor(
|
|
[vq_out_feats.shape[1]], dtype=torch.long, device=vq_out_feats.device
|
|
)
|
|
vq_out_feats, vq_out_length = self.upsample(vq_out_feats, vq_out_length)
|
|
# audio: (b, t)
|
|
audio, audio_length = self.acoustic_decoder(vq_out_feats, vq_out_length)
|
|
return audio
|
|
|
|
@torch.inference_mode()
|
|
def decode_one_token(
|
|
self, token: torch.Tensor, cache_dict: Dict[str, torch.Tensor], last_token: bool
|
|
):
|
|
"""Decode one single token to audio.
|
|
|
|
Args:
|
|
token: (B=1, nq, L=1)
|
|
Returns:
|
|
audio: (B=1, t)
|
|
"""
|
|
# token->latent->upsample, (naturally causal)
|
|
token = token.permute(1, 0, 2) # (B, nq, L) -> (nq, B, L)
|
|
vq_out_feats = self.rvq.decode_codes(token)
|
|
vq_out_feats = vq_out_feats.transpose(1, 2)
|
|
vq_out_length = torch.tensor(
|
|
[vq_out_feats.shape[1]], dtype=torch.long, device=vq_out_feats.device
|
|
)
|
|
vq_out_feats, vq_out_length = self.upsample(vq_out_feats, vq_out_length)
|
|
# acoustic decoder
|
|
up_conv_cache = cache_dict.get("up_conv_cache", None)
|
|
bb_conv_cache1 = cache_dict.get("bb_conv_cache1", None)
|
|
bb_conv_cache2 = cache_dict.get("bb_conv_cache2", None)
|
|
bb_kv_cache = cache_dict.get("bb_kv_cache", None)
|
|
is_cache = cache_dict.get("is_cache", None)
|
|
|
|
(
|
|
audio,
|
|
new_up_conv_cache,
|
|
new_bb_conv_cache1,
|
|
new_bb_conv_cache2,
|
|
new_bb_kv_cache,
|
|
new_is_cache,
|
|
) = self.acoustic_decoder.forward_chunk(
|
|
vq_out_feats,
|
|
up_conv_cache,
|
|
bb_conv_cache1,
|
|
bb_conv_cache2,
|
|
bb_kv_cache,
|
|
is_cache,
|
|
last_token,
|
|
)
|
|
|
|
new_cache_dict = {
|
|
"up_conv_cache": new_up_conv_cache,
|
|
"bb_conv_cache1": new_bb_conv_cache1,
|
|
"bb_conv_cache2": new_bb_conv_cache2,
|
|
"bb_kv_cache": new_bb_kv_cache,
|
|
"is_cache": new_is_cache,
|
|
}
|
|
return audio, new_cache_dict
|