CAMB-AI/MARS5-TTS
48076
1import math2from typing import Optional3 4import torch5import torch.nn as nn6import torch.nn.functional as F7from torch import Tensor8 9from .nn_future import (FNNSwiGLU, MistralTransformer, ModelArgs,10 RotatingBufferCache, SinePositionalEmbedding)11from .utils import construct_padding_mask, length_to_mask12 13LAYERNORM_EPS = 4e-514 15# ------------------------16# Code adapted from OpenAI guided diffusion repo17 18def timestep_embedding(timesteps, dim, max_period=10000, dtype=torch.float32):19 """20 Create sinusoidal timestep embeddings.21 :param timesteps: a 1-D Tensor of N indices, one per batch element.22 These may be fractional.23 :param dim: the dimension of the output.24 :param max_period: controls the minimum frequency of the embeddings.25 :return: an [N x dim] Tensor of positional embeddings.26 """27 half = dim // 228 freqs = torch.exp(29 -math.log(max_period) * torch.arange(start=0, end=half) / half30 ).to(device=timesteps.device)31 args = timesteps[:, None].float() * freqs[None]32 embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1).to(dtype)33 if dim % 2:34 embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)35 return embedding36 37 38# --------------------------------39# autoregressive codec language model40 41 42class CodecLM(nn.Module):43 44 def __init__(self, n_vocab, dim=1536, nhead=24, n_layers=26, n_spk_layers=2, dim_ff_scale=None, sliding_window=3000) -> None:45 super().__init__()46 47 if dim_ff_scale is None: hidden_dim = int(dim*4*(3/4))48 else: hidden_dim = int(dim*dim_ff_scale)49 50 self.cfg = ModelArgs(n_vocab, dim=dim, n_layers=n_layers, n_heads=nhead, n_kv_heads=nhead, hidden_dim=hidden_dim, sliding_window=sliding_window)51 self.ar = MistralTransformer(self.cfg)52 53 self.embed = nn.Embedding(n_vocab, dim)54 55 # --- spk embedding network56 dim_ff = int(dim*4*(3/4))57 self.pos_embedding = SinePositionalEmbedding(dim, scale=False, alpha=True)58 self.ref_chunked_emb = ChunkedEmbedding(1024 + 1, 8, dim) # add 1 for pad idx59 self.spk_identity_emb = nn.Embedding(1, dim)60 # define custom decoder61 encoder_layer = nn.TransformerEncoderLayer(dim, nhead, dim_ff,62 activation=FNNSwiGLU(dim, dim_ff), dropout=0,63 batch_first=True, norm_first=True, layer_norm_eps=LAYERNORM_EPS)64 encoder_layer.linear1 = nn.Identity()65 self.spk_encoder = nn.TransformerEncoder(encoder_layer, n_spk_layers, norm=nn.LayerNorm(dim, eps=LAYERNORM_EPS))66 # monkeypatch for broken copy.deepcopy of nn.Modules in nn.TransformerDecoder67 for l in self.spk_encoder.layers: l.activation = FNNSwiGLU(dim, dim_ff)68 69 70 @torch.inference_mode71 def get_spk_embedding(self, spk_reference, c_codes_lengths=None) -> Tensor:72 """ Gets speaker reference embeddings using `spk_reference` codes of shape (bs, seq_len, n_codebooks). """73 bs = spk_reference.shape[0]74 if bs != 1:75 raise AssertionError(f"Speaker embedding extraction only implemented using for bs=1 currently.")76 spk_seq = self.ref_chunked_emb(spk_reference) # (bs, sl, dim)77 spk_ref_emb = self.spk_identity_emb.weight[None].expand(bs, -1, -1) # (bs, 1, dim)78 79 spk_seq = torch.cat([spk_ref_emb, spk_seq], dim=1) # (bs, 1+sl, dim)80 # add pos encoding81 spk_seq = self.pos_embedding(spk_seq)82 # codebook goes from indices 0->1023, padding is idx 1024 (the 1025th entry)83 src_key_padding_mask = construct_padding_mask(spk_reference[:, :, 0], 1024) 84 src_key_padding_mask = torch.cat((85 # append a zero here since we DO want to attend to initial position.86 torch.zeros(src_key_padding_mask.shape[0], 1, dtype=bool, device=src_key_padding_mask.device), 87 src_key_padding_mask88 ), 89 dim=1)90 # pass through transformer91 res = self.spk_encoder(spk_seq, is_causal=False, src_key_padding_mask=src_key_padding_mask)[:, :1] # select first element -> now (bs, 1, dim).92 return res.squeeze(1)93 94 95 def forward(self, x: Tensor, x_padding_mask: Optional[Tensor] = None, spk_reference: Optional[Tensor] = None,96 cache: Optional[RotatingBufferCache] = None, counter: int = 0) -> Tensor:97 """ Inputs:98 - `x`: (bs, seq_len, vocab_size) 99 - `x_padding_mask`: (bs, seq_len) mask for each input, True for positions to *ignore*, False otherwise.100 Note that since this is an autoregressive model, this doesn't actually matter for infernece, so it is ignored at inference. 101 - `spk_reference`: (bs, seq_len, n_codebooks) corresponding to the speaker reference to clone from.102 - `cache` and `counter`: used for kv caching, optional.103 104 Returns `x` of same shape (bs, seq_len, dim)105 """106 x = self.embed(x)107 108 # --- speaker reference/embedding109 if spk_reference is not None:110 # compute ref111 bs = spk_reference.shape[0]112 spk_seq = self.ref_chunked_emb(spk_reference) # (bs, sl, dim)113 spk_ref_emb = self.spk_identity_emb.weight[None].expand(bs, -1, -1) # (bs, 1, dim)114 115 spk_seq = torch.cat([spk_ref_emb, spk_seq], dim=1) # (bs, 1+sl, dim)116 # add pos encoding117 spk_seq = self.pos_embedding(spk_seq)118 # codebook goes from indices 0->1023, padding is idx 1024 (the 1025th entry)119 src_key_padding_mask = construct_padding_mask(spk_reference[:, :, 0], 1024) 120 src_key_padding_mask = torch.cat((121 # append a zero here since we DO want to attend to initial position.122 torch.zeros(src_key_padding_mask.shape[0], 1, dtype=bool, device=src_key_padding_mask.device), 123 src_key_padding_mask124 ), 125 dim=1)126 # pass through transformer127 res = self.spk_encoder(spk_seq, is_causal=False, src_key_padding_mask=src_key_padding_mask)[:, :1] # select first element -> now (bs, 1, dim).128 129 x = torch.cat([res, x], dim=1)130 131 positions = torch.arange(0, x.shape[1], device=x.device, dtype=torch.long)132 if cache is not None and counter != 1:133 # using only the last token to predict the next one134 x = x[:,-1,:].unsqueeze(1)135 positions = positions[-1:]136 137 x = self.ar(x, positions, cache) # (bs, seq_len, vocab)138 if spk_reference is not None and (cache is None or counter == 1):139 x = x[:, 1:] # strip out the first output token corresponding to the speaker embedding token.140 141 return x142 143 144# -------------------------145# residual discrete diffusion model146 147class ChunkedEmbedding(nn.Module):148 149 def __init__(self, codebook_size: int, n_quantizer: int, dim: int) -> None:150 super().__init__()151 assert dim % n_quantizer == 0, f"ChunkedEmbedding output dim ({dim}) must be divisible by n_quant {n_quantizer}"152 self.embs = nn.ModuleList([nn.Embedding(codebook_size, dim//n_quantizer) for _ in range(n_quantizer)])153 154 def forward(self, x: Tensor) -> Tensor:155 """ Embeds each codebook index in `x` (bs, seq_len, n_quantizer) to an embedding vector, concatenating results.156 Returns output of shape (bs, seq_len, dim)157 """158 y = torch.cat([self.embs[i](x[..., i]) for i in range(x.shape[-1])], dim=-1)159 return y160 161 162 163class ResidualTransformer(nn.Module):164 165 def __init__(self, n_text_vocab, n_quant=1024, dim=1024, nhead=16, 166 enc_layers=8, dec_layers=16, n_spk_layers=3,167 c_quant_levels=8, pred_quant_levels=8, 168 t_emb_dim=1024, norm_first=True, p_cond_drop=0.1, dropout=0) -> None:169 super().__init__()170 171 self.cond_pos_embedding = SinePositionalEmbedding(dim, scale=False, alpha=True)172 self.pos_embedding = SinePositionalEmbedding(dim, scale=False, alpha=True)173 174 # *4 from heuristic, *2/3 from swiglu, since there are 3 linear matrices not 2.175 # so we must keep # params the same.176 dim_ff = int(dim*4*(3/4))177 178 # define custom encoder179 encoder_layer = nn.TransformerEncoderLayer(dim, nhead, dim_ff,180 activation=FNNSwiGLU(dim, dim_ff), dropout=dropout,181 batch_first=True, norm_first=norm_first, layer_norm_eps=LAYERNORM_EPS)182 encoder_layer.linear1 = nn.Identity()183 encoder = nn.TransformerEncoder(encoder_layer, enc_layers, norm=nn.LayerNorm(dim, eps=LAYERNORM_EPS) if norm_first else None)184 185 # define custom decoder186 decoder_layer = nn.TransformerDecoderLayer(dim, nhead, dim_ff,187 activation=FNNSwiGLU(dim, dim_ff), dropout=dropout,188 batch_first=True, norm_first=norm_first, layer_norm_eps=LAYERNORM_EPS)189 decoder_layer.linear1 = nn.Identity()190 decoder = nn.TransformerDecoder(decoder_layer, dec_layers, norm=nn.LayerNorm(dim, eps=LAYERNORM_EPS) if norm_first else None)191 192 # monkeypatch for broken copy.deepcopy of nn.Modules in nn.TransformerDecoder193 for l in decoder.layers: l.activation = FNNSwiGLU(dim, dim_ff)194 195 self.tfm = nn.Transformer(dim, nhead, dim_feedforward=dim_ff, batch_first=True, 196 norm_first=norm_first,197 num_encoder_layers=enc_layers,198 num_decoder_layers=dec_layers,199 custom_encoder=encoder,200 custom_decoder=decoder,201 layer_norm_eps=LAYERNORM_EPS,202 dropout=dropout203 )204 # Timestep embedding network205 self.t_emb_dim = t_emb_dim206 self.timestep_encoder_emb = nn.Sequential(207 nn.Linear(t_emb_dim, dim),208 nn.SiLU(),209 nn.Linear(dim, dim)210 )211 self.timestep_decoder_emb = nn.Sequential(212 nn.Linear(t_emb_dim, dim),213 nn.SiLU(),214 nn.Linear(dim, dim)215 )216 217 self.text_embed = nn.Embedding(n_text_vocab, dim)218 219 ## ----> reference / conditioning encoder:220 self.ref_embedder = ChunkedEmbedding(n_quant, c_quant_levels, dim)221 self.ref_pos_embedding = SinePositionalEmbedding(dim, scale=False, alpha=True)222 self.spk_identity_emb = nn.Embedding(1, dim)223 spk_encoder_layer = nn.TransformerEncoderLayer(dim, nhead, dim_ff,224 activation=FNNSwiGLU(dim, dim_ff), dropout=dropout,225 batch_first=True, norm_first=True, layer_norm_eps=LAYERNORM_EPS)226 spk_encoder_layer.linear1 = nn.Identity()227 self.spk_encoder = nn.TransformerEncoder(spk_encoder_layer, n_spk_layers, norm=nn.LayerNorm(dim, eps=LAYERNORM_EPS))228 # monkeypatch for broken copy.deepcopy of nn.Modules in nn.TransformerDecoder229 for l in self.spk_encoder.layers: l.activation = FNNSwiGLU(dim, dim_ff)230 # ----> end speaker encoder network231 232 # self.residual_encoder = nn.Embedding(n_quant, dim) # only encode first quantization level of decoder input.233 self.residual_encoder = ChunkedEmbedding(n_quant, c_quant_levels, dim)234 235 self.residual_decoder = nn.ModuleList([236 nn.Sequential(237 nn.LayerNorm(dim),238 nn.Linear(dim, n_quant)239 ) for i in range(pred_quant_levels)240 ])241 self.n_quantizer = pred_quant_levels242 self.p_cond_drop = p_cond_drop243 244 245 @torch.inference_mode246 def get_spk_embedding(self, c_codes, c_codes_length) -> Tensor:247 """ Obtain speaker embedding vectors using `c_codes` from reference encodec sequences, and `c_codes_length` of lengths for each sequence """248 bs = c_codes.shape[0]249 spk_seq = self.ref_embedder(c_codes) # (bs, sl, dim)250 spk_ref_emb = self.spk_identity_emb.weight[None].expand(bs, -1, -1) # (bs, 1, dim)251 spk_seq = torch.cat([spk_ref_emb, spk_seq], dim=1) # (bs, 1+sl, dim)252 # add pos encoding253 spk_seq = self.ref_pos_embedding(spk_seq)254 255 # add 1 to c_codes_length to account for the fact that we concatenate the spk_ref_emb to it. 256 src_key_padding_mask = length_to_mask(c_codes_length+1, torch.zeros_like(c_codes_length), max_len=spk_seq.shape[1])257 src_key_padding_mask = src_key_padding_mask.to(dtype=torch.bool, device=spk_seq.device)258 259 # pass through transformer260 res = self.spk_encoder(spk_seq, is_causal=False, src_key_padding_mask=src_key_padding_mask)[:, :1] # select first element -> now (bs, 1, dim).261 return res.squeeze(1)262 263 264 def forward(self, c_text: Tensor, c_codes: Tensor, c_texts_length: Tensor, c_codes_length: Tensor, 265 x: Tensor, x_padding_mask: Tensor, t: Tensor, drop_cond=False):266 """ Input:267 - `c_text`: (bs, seq_len1) the prompt text (BPE encoded)268 - `c_codes`: (bs, seq_len2, n_quant) the full tokenized codes of the reference speech269 - `c_texts_length`: (bs, ) the length of the codes in the text prompt270 - `c_codes_length`: (bs, ) the length of the prompt acoustic token codes in `c_codes`.271 - `x`: (bs, seq_len3) L0 residual codes272 - `x`: (bs, seq_len3, n_quant) L0 residual codes273 - `x_padding_mask`: (bs, seq_len3) masking for residual codes274 - `t`: (bs) timestep275 - `drop_cond`: bool, whether or not to forcibly drop the conditioning information.276 Returns:277 - outs: (bs, seq_len, n_quantizer, codebook_size)278 """279 280 c_text = self.text_embed(c_text) # (bs, seq_len1, dim)281 282 ## ----> reference / conditioning encoder:283 bs = c_codes.shape[0]284 285 286 if self.training:287 zero_cond_inds = torch.rand_like(t, dtype=c_text.dtype) < self.p_cond_drop288 else:289 # never randomly zero when in eval mode290 zero_cond_inds = torch.zeros_like(t, dtype=torch.bool)291 if drop_cond:292 # force drop conditioning293 zero_cond_inds = torch.ones_like(t, dtype=torch.bool)294 295 c_codes_length[zero_cond_inds] = 0296 c_codes[zero_cond_inds] = 1024297 298 spk_seq = self.ref_embedder(c_codes) # (bs, sl, dim)299 spk_ref_emb = self.spk_identity_emb.weight[None].expand(bs, -1, -1) # (bs, 1, dim)300 spk_seq = torch.cat([spk_ref_emb, spk_seq], dim=1) # (bs, 1+sl, dim)301 # add pos encoding302 spk_seq = self.ref_pos_embedding(spk_seq)303 304 # add 1 to c_codes_length to account for the fact that we concatenate the spk_ref_emb to it. 305 src_key_padding_mask = length_to_mask(c_codes_length+1, torch.zeros_like(c_codes_length), max_len=spk_seq.shape[1])306 src_key_padding_mask = src_key_padding_mask.to(dtype=torch.bool, device=spk_seq.device)307 308 # pass through transformer309 res = self.spk_encoder(spk_seq, is_causal=False, src_key_padding_mask=src_key_padding_mask)[:, :1] # select first element -> now (bs, 1, dim).310 c_codes = res # (bs, 1, dim)311 c_codes_lengths_extract = torch.ones_like(c_codes_length) # manually override all the code lengths to equal 1, since we only have 1 spk embedding. 312 ## ----> end reference / conditioning encoder:313 314 ## ----> timestep embeddings and parsing315 t_emb = timestep_embedding(t, self.t_emb_dim, dtype=c_text.dtype)316 t_emb_encoder = self.timestep_encoder_emb(t_emb) # (bs, t_dim)317 t_emb_decoder = self.timestep_decoder_emb(t_emb)318 319 ## ----> concatenating text/phone inputs and implicit speaker embedding. 320 c_phones_unpacked = nn.utils.rnn.unpad_sequence(c_text, c_texts_length.cpu(), batch_first=True)321 c_codes_unpacked = nn.utils.rnn.unpad_sequence(c_codes, c_codes_lengths_extract.cpu(), batch_first=True)322 # >>> Concat [speaker codes, text codes]323 assert all(b.shape[0] == 1 for b in c_codes_unpacked)324 c_joined = [torch.cat((b, a), dim=0) for a, b in zip(c_phones_unpacked, c_codes_unpacked)]325 326 c = nn.utils.rnn.pad_sequence(c_joined, batch_first=True)327 c_joined_lengths = torch.tensor([p.shape[0] for p in c_joined], device=c.device, dtype=torch.long)328 c_padding_mask = length_to_mask(c_joined_lengths, torch.zeros_like(c_joined_lengths))329 c = self.cond_pos_embedding(c)330 331 ## Format input:332 x = self.residual_encoder(x) # (bs, seq_len3, dim)333 334 x = self.pos_embedding(x)335 336 x = x + t_emb_decoder[:, None]337 c = c + t_emb_encoder[:, None]338 ## Perform prediction:339 output = self.tfm(c, x, src_key_padding_mask=c_padding_mask, 340 tgt_key_padding_mask=x_padding_mask,341 memory_key_padding_mask=c_padding_mask) # (bs, seq_len, dim)342 outs = torch.stack([self.residual_decoder[i](output) for i in range(self.n_quantizer)], dim=-1) # (bs, seq_len, logit_dim, n_quant)343 return outs344 345