import random import numpy as np import logging import argparse, copy import torch import torch.nn as nn import torch.nn.functional as F from torchmetrics.classification import MulticlassAccuracy from .modules.utils import make_pad_mask from .modules.embedding import SinePositionalEmbedding, TokenEmbedding from .modules.transformer import ( LayerNorm, TransformerEncoder, TransformerEncoderLayer, ) from .codebooks_patterns import DelayedPatternProvider def top_k_top_p_filtering( logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1 ): """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering Args: logits: logits distribution shape (batch size, vocabulary size) if top_k > 0: keep only top k tokens with highest probability (top-k filtering). if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering). Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751) Make sure we keep at least min_tokens_to_keep per batch example in the output From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317 """ if top_k > 0: top_k = min( max(top_k, min_tokens_to_keep), logits.size(-1) ) # Safety check # Remove all tokens with a probability less than the last token of the top-k indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None] logits[indices_to_remove] = filter_value if top_p < 1.0: sorted_logits, sorted_indices = torch.sort(logits, descending=True) cumulative_probs = torch.cumsum( F.softmax(sorted_logits, dim=-1), dim=-1 ) # Remove tokens with cumulative probability above the threshold (token with 0 are kept) sorted_indices_to_remove = cumulative_probs > top_p if min_tokens_to_keep > 1: # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below) sorted_indices_to_remove[..., :min_tokens_to_keep] = 0 # Shift the indices to the right to keep also the first token above the threshold sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[ ..., :-1 ].clone() sorted_indices_to_remove[..., 0] = 0 # scatter sorted tensors to original indexing indices_to_remove = sorted_indices_to_remove.scatter( 1, sorted_indices, sorted_indices_to_remove ) logits[indices_to_remove] = filter_value return logits def topk_sampling(logits, top_k=10, top_p=1.0, temperature=1.0): # temperature: (`optional`) float # The value used to module the next token probabilities. Must be strictly positive. Default to 1.0. # top_k: (`optional`) int # The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50. # top_p: (`optional`) float # The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1. # Temperature (higher temperature => more likely to sample low probability tokens) if temperature != 1.0: logits = logits / temperature # Top-p/top-k filtering logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p) # Sample token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1) return token class VoiceCraft(nn.Module): def __init__(self, args): super().__init__() self.args = copy.copy(args) self.pattern = DelayedPatternProvider(n_q=self.args.n_codebooks) if not getattr(self.args, "special_first", False): self.args.special_first = 0 if not getattr(self.args, "n_special", False): self.args.n_special = 3 self.args.eos = getattr(self.args, "eos", -1) self.eog = nn.Parameter(torch.full((self.args.n_codebooks, 1), self.args.eog, dtype=torch.long), requires_grad=False) # [K 1] if self.args.eos > 0: assert self.args.eos != self.args.audio_pad_token and self.args.eos != self.args.empty_token, self.args.eos self.eos = nn.Parameter(torch.full((self.args.n_codebooks, 1), self.args.eos, dtype=torch.long), requires_grad=False) # [K 1] if type(self.args.audio_vocab_size) == str: self.args.audio_vocab_size = eval(self.args.audio_vocab_size) self.n_text_tokens = self.args.text_vocab_size + 1 assert self.args.text_pad_token == self.args.text_vocab_size, f"self.args.text_vocab_size: {self.args.text_vocab_size}, self.args.text_pad_token: {self.args.text_pad_token}" self.n_audio_tokens = [self.args.audio_vocab_size + self.args.n_special] * self.args.n_codebooks # special tokens: empty token, EOG token, audio pad token assert self.args.audio_vocab_size == self.args.empty_token, self.args.empty_token assert self.args.eog == self.args.audio_vocab_size + 1, self.args.eog assert self.args.audio_pad_token == self.args.audio_vocab_size + 2, self.args.audio_pad_token self.text_embedding = TokenEmbedding( dim_model=self.args.d_model, vocab_size=self.n_text_tokens, dropout=self.args.text_embedding_dropout ) self.audio_embedding = nn.ModuleList( [ TokenEmbedding( dim_model=self.args.audio_embedding_dim, vocab_size=self.n_audio_tokens[k], dropout=self.args.audio_embedding_dropout ) for k in range(self.args.n_codebooks) ] ) self.mask_embedding = nn.Parameter(torch.randn(self.args.max_n_spans, self.args.d_model), requires_grad=True) self.text_positional_embedding = SinePositionalEmbedding( self.args.d_model, dropout=self.args.text_positional_embedding_dropout, scale=False, alpha=True, # learnable scaler, scale the volume of positional embedding ) self.audio_positional_embedding = SinePositionalEmbedding( self.args.d_model, dropout=self.args.audio_positional_embedding_dropout, scale=False, alpha=True, # learnable scaler, scale the volume of positional embedding ) dec_layer = TransformerEncoderLayer( self.args.d_model, self.args.nhead, dim_feedforward=self.args.d_model * 4, dropout=self.args.trm_dropout, batch_first=True, norm_first=True, layer_norm_cls=LayerNorm ) self.decoder = TransformerEncoder( dec_layer, num_layers=self.args.num_decoder_layers, norm=LayerNorm(self.args.d_model), ) self.predict_layer = nn.ModuleList( [ nn.Sequential(nn.Linear(self.args.d_model, self.args.audio_vocab_size//2), nn.GELU(), nn.Linear(self.args.audio_vocab_size//2, self.n_audio_tokens[k])) for k in range(self.args.n_codebooks) ] ) self.accuracy_metrics = nn.ModuleList( [MulticlassAccuracy( self.n_audio_tokens[k], top_k=10, average="micro", multidim_average="global", ignore_index=None, ) for k in range(self.args.n_codebooks)] ) def prepare_mask_intervals(self, y_lens): mask_intervals = [] non_mask_intervals = [] for i, y_len in enumerate(y_lens): if self.args.mask_sample_dist == "uniform": n_spans = random.choice(range(1, self.args.max_n_spans+1)) elif "poisson" in self.args.mask_sample_dist.lower(): param = float(self.args.mask_sample_dist[len("poisson"):]) poisson_sample = torch.poisson(torch.tensor([param])) n_spans = int(poisson_sample.clamp(1, self.args.max_n_spans).item()) starts = random.sample(range(1, y_len-1-self.args.mask_len_min), n_spans) starts = sorted(starts) for j in range(len(starts)-1, 0, -1): if starts[j] - starts[j-1] < self.args.min_gap: del starts[j] # If elements are too close, delete the later one assert len(starts) > 0, f"there is no masked span left, y_len: {y_len}, sampled n_spans: {n_spans}" temp_starts = starts + [y_len] gaps = [temp_starts[j+1] - temp_starts[j] for j in range(len(temp_starts)-1)] ends = [] for j, (start, gap) in enumerate(zip(starts, gaps)): mask_len = random.randint(self.args.mask_len_min, self.args.mask_len_max) # if mask_len > gap * self.args.max_mask_portion: # make sure the masks are not overlapping with each other if mask_len > gap - 1: # make sure the masks are not overlapping with each other # temp_mask_start = int(0.6*gap*self.args.max_mask_portion) # temp_mask_end = int(gap*self.args.max_mask_portion) temp_mask_start = 1 temp_mask_end = gap - 1 mask_len = random.randint(temp_mask_start, temp_mask_end) ends.append(start + mask_len) mask_intervals.append([(s,e) for s,e in zip(starts, ends)]) non_mask_intervals.append([(ns,ne) for ns, ne in zip([0]+ends, starts+[y_len])]) return mask_intervals, non_mask_intervals def rearrange(self, y, non_mask_intervals, mask_intervals): reduced_eog = getattr(self.args, "reduced_eog", 0) rearranged_y = [] for i in range(len(y)): if self.args.eos > 0: assert reduced_eog cur_y = [y[i, :, item[0]: item[1]] for item in non_mask_intervals[i][:-1]] + [torch.cat([y[i, :, non_mask_intervals[i][-1][0]: non_mask_intervals[i][-1][1]], self.eos], dim=-1)] + [torch.cat([y[i, :, item[0]: item[1]], self.eog], dim=-1) for item in mask_intervals[i]] # only insert eog to the last non-mask-interval, which is when the utterance actual ends else: if reduced_eog: cur_y = [y[i, :, item[0]: item[1]] for item in non_mask_intervals[i][:-1]] + [torch.cat([y[i, :, non_mask_intervals[i][-1][0]: non_mask_intervals[i][-1][1]], self.eog], dim=-1)] + [torch.cat([y[i, :, item[0]: item[1]], self.eog], dim=-1) for item in mask_intervals[i]] # only insert eog to the last non-mask-interval, which is when the utterance actual ends else: cur_y = [torch.cat([y[i, :, item[0]: item[1]], self.eog], dim=-1) for item in non_mask_intervals[i]] + [torch.cat([y[i, :, item[0]: item[1]], self.eog], dim=-1) for item in mask_intervals[i]] # eog is added to each section TODO this is not correct, I should add eog to non_mask_intervals if that segment is not the ending segment (as there is no way for the model to predict eog for those segments, and this will do harm to tts experiment, where the model randomly output eog for the first segment) rearranged_y.append(cur_y) return rearranged_y def shift(self, rearranged_y): shifted_y = [] patterns = [] for i in range(len(rearranged_y)): cur_patterns = [self.pattern.get_pattern(cur_y.shape[1]) for cur_y in rearranged_y[i]] out = [cur_pattern.build_pattern_sequence(z=cur_y.unsqueeze(0).contiguous(), special_token=self.args.empty_token, keep_only_valid_steps=False) for cur_pattern, cur_y in zip(cur_patterns, rearranged_y[i])] shifted_y.append([item[0].squeeze(0) for item in out]) # the first item is values, later two are indexes and mask patterns.append(cur_patterns) return shifted_y, patterns def insert_mask(self, shifted_y): inserted_y = [] mask_position = [] mask_value = [] for i in range(len(shifted_y)): num_masks = (len(shifted_y[i]) - 1) // 2 assert num_masks == (len(shifted_y[i]) - 1) / 2, len(shifted_y[i]) emb_inds = list(range(self.args.max_n_spans)) if self.args.shuffle_mask_embedding: random.shuffle(emb_inds) emb_inds_use = emb_inds[:num_masks] emb_inds_use = emb_inds_use + emb_inds_use mask_value.append(emb_inds_use) cur_inserted_y = [] cur_mask_position = [] for j in range(len(shifted_y[i])-1): cur_inserted_y.append(shifted_y[i][j]) cur_mask_position.append(sum([item.shape[1] for item in cur_inserted_y])) # each item is of shape [K S], so take shape[1] cur_inserted_y.append(self.eog) # insert mask token of shape [K, 1], BUT we are actually using the eog token as a place holder here, as the real mask will be inserted in embed_y function cur_inserted_y.append(shifted_y[i][-1]) inserted_y.append(cur_inserted_y) mask_position.append(cur_mask_position) return inserted_y, mask_position, mask_value def cat_y(self, inserted_y, mask_position, y_lens): reduced_eog = getattr(self.args, "reduced_eog", 0) cated_y = [] new_y_lens = [] for i in range(len(inserted_y)): cur_cated_y = torch.cat(inserted_y[i], dim=1) #[K S] cur_cated_y = cur_cated_y.transpose(1,0) # [S K] cur_cated_y_len = cur_cated_y.shape[0] if reduced_eog: assert cur_cated_y_len == y_lens[i] + len(mask_position[i]) + (len(mask_position[i]) + 1) * self.args.n_codebooks + (len(mask_position[i])/2 + 1), f"cur_cated_y_len == {cur_cated_y_len}, but it should be y_lens[i] ({y_lens[i]}) + len(mask_position[i]) ({len(mask_position[i])}) + (len(mask_position[i]) + 1) * self.args.n_codebooks ({(len(mask_position[i]) + 1) * self.args.n_codebooks}) + (len(mask_position[i])/2 + 1) ({len(mask_position[i])/2 + 1})={y_lens[i] + len(mask_position[i]) + (len(mask_position[i]) + 1) * self.args.n_codebooks + (len(mask_position[i])/2 + 1)}" else: assert cur_cated_y_len == y_lens[i] + len(mask_position[i]) + (len(mask_position[i]) + 1) * self.args.n_codebooks + (len(mask_position[i]) + 1), f"cur_cated_y_len == {cur_cated_y_len}, but it should be y_lens[i] ({y_lens[i]}) + len(mask_position[i]) ({len(mask_position[i])}) + (len(mask_position[i]) + 1) * self.args.n_codebooks ({(len(mask_position[i]) + 1) * self.args.n_codebooks}) + (len(mask_position[i]) + 1) ({len(mask_position[i]) + 1})" # the last term represent the inserted eog token, originally it's inserted at the end of every token, but this is wrong new_y_lens.append(cur_cated_y_len) cated_y.append(cur_cated_y) cated_y = torch.nn.utils.rnn.pad_sequence(cated_y, batch_first=False, padding_value=self.args.audio_pad_token) assert cated_y.shape == torch.Size([max(new_y_lens),len(inserted_y), self.args.n_codebooks]), f"cated_y.shape: {cated_y.shape}, but it should be {torch.Size([max(new_y_lens,len(inserted_y), self.args.n_codebooks)])}" cated_y = cated_y.permute(2,0,1) # [T,B,K]->[K,T,B] assert cated_y.shape[0] == self.args.n_codebooks, cated_y.shape return cated_y, torch.LongTensor(new_y_lens).to(cated_y.device) def embed_y(self, cated_y, mask_position, mask_value): embedded_y = torch.stack([self.audio_embedding[k](cated_y[k]) for k in range(self.args.n_codebooks)], dim=0) # [K, T, B, D] assert embedded_y.shape[0] == self.args.n_codebooks, embedded_y.shape assert embedded_y.shape[-1] == self.args.d_model, embedded_y.shape embedded_y = embedded_y.sum(dim=0) # [K,T,B,D]->[T,B,D] embedded_y = embedded_y.transpose(1,0) # [T,B,D]->[B,T,D] for i in range(len(embedded_y)): if len(mask_position[i]) > 0: embedded_y[i, mask_position[i]] = self.mask_embedding[mask_value[i]] return embedded_y def prepare_input_target(self, y, y_lens): # rearrange y # assume y shape: [B T K], K is n_codebooks assert y.shape[1] == self.args.n_codebooks, y.shape # sample mask_intervals mask_intervals, non_mask_intervals = self.prepare_mask_intervals(y_lens) # need to have EOG in each section (SOG will be generated by the pattern class) # but mask can be inserted later after we have shifted the input # y could be rearranged in this way: # [ # [tensor[4, 12], tensor[4, 45], tensor[4, 102], tensor[4, 32]], tensor[4, 22]], # [tensor[4, 44], tensor[4, 56], tensor[4, 19]], # ... # ] # for the first list of tensors (4 tensors), first 3 tensors are non_masked part, last 2 are masked part. # NOTE #non_masked_part = #masked_part + 1 # NOTE *these are also the targets* # added eog at the end of each segment (masked segment and unmasked segment) rearranged_y = self.rearrange(y, non_mask_intervals, mask_intervals) targets = rearranged_y # each element in each sample is of shape [K T] assert targets[0][0].shape[0] == self.args.n_codebooks, targets[0][0].shape # next we need to apply pattern shifting to each tensor, after which, we'll replace the starting tokens of each section with a token that's different from the special padding token # [[5, 1, 2, 3, 4, 5, 5], # [5, 5, 1, 2, 3, 4, 5], # [5, 5, 5, 1, 2, 3, 4]] shifted_y, patterns = self.shift(rearranged_y) # each element [K S] assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape[0] # then, insert mask token at the intersection of each tensor (we want to decide the arrangement of the mask (shuffle or not)), we better have a separate nn.embedding for it # we also need to record the position of the inserted mask inserted_y, mask_position, mask_value = self.insert_mask(shifted_y) assert inserted_y[0][0].shape[0] == self.args.n_codebooks, inserted_y[0][0].shape[0] assert inserted_y[0][1].shape == torch.Size((self.args.n_codebooks, 1)), f"this should be a mask, so should have shape {(self.args.n_codebooks, 1)}, but it's {inserted_y[0][1].shape}" # then concat tensors that belong to the same sample (in order) then get the length of each sample, and then stack them in batch dimension, pad them with pad_token cated_y, new_y_lens = self.cat_y(inserted_y, mask_position, y_lens) # KTB assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], len(inserted_y))) # embed remember to separately embed the mask tokens embedded_y = self.embed_y(cated_y, mask_position, mask_value) #BTD assert embedded_y.shape[1:] == torch.Size((max(new_y_lens), self.args.d_model)), embedded_y.shape # positional embedding y_input = self.audio_positional_embedding(embedded_y) # make attention mask and padding mask y_padding_mask = make_pad_mask(new_y_lens).to(y.device) y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y_padding_mask.device) return y_input, new_y_lens, targets, y_padding_mask, y_attention_mask, mask_position, patterns def remove_mask(self, logits, mask_position, new_y_lens): # logits: [B K S card] logits_use = [] for i in range(len(logits)): non_mask_positions = [-1] + mask_position[i] + [new_y_lens[i]] non_mask_intervals = [[non_mask_positions[i]+1, non_mask_positions[i+1]] for i in range(len(non_mask_positions)-1)] cur_logits_use = [logits[i, :, l:r] for l,r in non_mask_intervals] logits_use.append(cur_logits_use) return logits_use def revert_pattern(self, patterns, logits_use): logits_final = [] logit_masks = [] for i in range(len(logits_use)): cur_logits = [ item.unsqueeze(0).permute(0, 3, 1, 2).contiguous() for item in logits_use[i] ] # each item is of shape [1 K S card] [1 card K S] cur_logits_final = [ cur_pattern.revert_pattern_logits( item, 0, keep_only_valid_steps=False ) for cur_pattern, item in zip(patterns[i], cur_logits) ] # if input output order doesn't match, this step will give an error cur_logits_final_ret = [item[0].permute(0,2,3,1).squeeze(0) for item in cur_logits_final] # each element is of shape [K,T,card] logits_final.append(cur_logits_final_ret) logit_masks.append([item[2] for item in cur_logits_final]) return logits_final, logit_masks def dec_forward( self, x_input, x_lens, x_attention_mask, x_padding_mask, y_input, new_y_lens, y_attention_mask, y_padding_mask, past=None, last_3_tokens=False ): x_attn_mask = F.pad( x_attention_mask, (0, new_y_lens.max()), value=True, ) # x attn to all x, doesn't attn to any y, this follow figure 3 of the valle paper y_attn_mask = F.pad( y_attention_mask, (x_lens.max(), 0), # y is padded at the front value=False, ) # y attn to all x, for y itself use lower triangle mask to ensure autoregressive xy_attn_mask = torch.concat([x_attn_mask, y_attn_mask], dim=0) # merge key padding and attention masks bsz, src_len = x_input.shape[0], x_lens.max() + new_y_lens.max() xy_padding_mask = torch.concat([x_padding_mask, y_padding_mask], dim=1) _xy_padding_mask = ( xy_padding_mask.view(bsz, 1, 1, src_len) .expand(-1, self.args.nhead, -1, -1) .reshape(bsz * self.args.nhead, 1, src_len) ) xy_attn_mask = xy_attn_mask.logical_or(_xy_padding_mask) new_attn_mask = torch.zeros_like(xy_attn_mask) new_attn_mask.masked_fill_(xy_attn_mask, float("-inf")) xy_attn_mask = new_attn_mask xy_input = torch.cat([x_input, y_input], dim=1) if past == None: # do not use kvcache out, _ = self.decoder((xy_input, None), mask=xy_attn_mask) return out[:, x_lens.max():], None else: # use kvcache if past.ndim > 3: # uses kvcache, only need to pass the last tokens, this doesn't work with multi-span speech editing yet if last_3_tokens: xy_input = xy_input[:, -3:] xy_attn_mask = xy_attn_mask[:, -3:] else: xy_input = xy_input[:, -1:] xy_attn_mask = xy_attn_mask[:, -1:] out, present = self.decoder((xy_input, None), mask=xy_attn_mask, past=past) if isinstance(out, tuple): # get rid of stage_embedding out = out[0] if out.shape[1] > x_lens.max(): # the first pass, not kvcache yet return out[:, x_lens.max():], present else: # used kvcache return out, present def forward(self, batch): """ Args: x: A 2-D tensor of shape (N, S). x_lens: A 1-D tensor of shape (N,). It contains the number of tokens in `x` before padding. y: A 3-D tensor of shape (N, K, T). where K is the number of codebooks y_lens: A 1-D tensor of shape (N,). It contains the number of tokens in `x` before padding. """ x, x_lens, y, y_lens = batch["x"], batch["x_lens"], batch["y"], batch["y_lens"] x = x[:, :x_lens.max()] # this deal with gradient accumulation, where x_lens.max() might not be longer than the length of the current slice of x y = y[:, :y_lens.max()] assert x.ndim == 2, x.shape assert x_lens.ndim == 1, x_lens.shape assert y.ndim == 3 and y.shape[1] == self.args.n_codebooks, y.shape assert y_lens.ndim == 1, y_lens.shape # makes attention mask and padding mask for x x_padding_mask = make_pad_mask(x_lens).to(x.device) x_attention_mask = torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1).bool().to(x_padding_mask.device) x_input = self.text_embedding(x) x_input = self.text_positional_embedding(x_input) y_input, new_y_lens, targets, y_padding_mask, y_attention_mask, mask_position, patterns = self.prepare_input_target(y, y_lens) y_out = self.dec_forward( x_input, x_lens, x_attention_mask, x_padding_mask, y_input, new_y_lens, y_attention_mask, y_padding_mask ) y_out = y_out[0] # no kv-caching during training assert y_out.shape == y_input.shape, f"y_out.shape: {y_out.shape}, y_input.shape: {y_input.shape}" # [B S D] logits = torch.stack([self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1) # [B K S card] # take out the mask token (using mask_position and new_y_lens) and revert (using function provided by self.pattern) assert logits.shape[1] == self.args.n_codebooks and logits.shape[3] == self.n_audio_tokens[0], logits.shape logits_use = self.remove_mask(logits, mask_position, new_y_lens) # revert the pattern shift for each logits section in each sample logits_final, logit_masks = self.revert_pattern(patterns, logits_use) assert logits_final[0][0].shape[0] == self.args.n_codebooks and logits_final[0][0].shape[2] == self.n_audio_tokens[0], f"it is: {logits_final[0][0].shape}, but should be [K, T, card]" # testing sample_to_test = 0 assert len(logits_final[sample_to_test]) == len(targets[sample_to_test]), f"{len(logits_final[sample_to_test])}, {len(targets[sample_to_test])}" temp = sum([logits_final[sample_to_test][i].shape[:-1] != targets[sample_to_test][i].shape for i in range(len(targets[sample_to_test]))]) assert temp == 0, f"none equal positions: {temp}, total number of elements: {len(targets[sample_to_test])}" logit_masked = sum([(item==False).any() for cur_mask in logit_masks for item in cur_mask]) assert logit_masked == 0, logit_masks logits = torch.cat([torch.cat(item, dim=1) for item in logits_final], dim=1) # [K, T1+T2+T3+..., card] targets = torch.cat([torch.cat(item, dim=1) for item in targets], dim=1) # [K, T1+T2+T3+...] assert targets.shape[0] == logits.shape[0], f"{targets.shape}, {logits.shape}" loss = [] ntokens = [] top10acc = [] for k, (logit, target) in enumerate(zip(logits, targets)): loss.append(F.cross_entropy(logit, target, reduction='mean')) top10acc.append(self.accuracy_metrics[k](logit.detach(), target)) ntokens.append(len(logit)) all_ntokens = sum(ntokens) if self.args.codebook_weight != None: codebook_weight = eval(self.args.codebook_weight) else: codebook_weight = [1.] * self.args.n_codebooks loss = sum([l*nt*cw for l, nt, cw in zip(loss, ntokens, codebook_weight)]) top10acc_by_codebook = [t10a*nt for t10a, nt in zip(top10acc, ntokens)] top10acc = sum(top10acc_by_codebook) ntokens = torch.tensor(all_ntokens).to(logits.device) return { "loss": loss, "top10acc": top10acc, "top10acc_by_codebook": top10acc_by_codebook, "effective_ntoken": ntokens, } def inference( self, x: torch.Tensor, x_lens: torch.Tensor, y: torch.Tensor, mask_interval: list[torch.Tensor], top_k: int=-100, top_p: float=1.0, temperature: float=1.0, stop_repetition: int=-1, kvcache: int=1, silence_tokens: list[int]=[1388,1898,131], ) -> torch.Tensor: """ Args: x: A 2-D tensor of shape (1, L). x_lens: A 1-D tensor of shape (1,). It contains the number of tokens in `x` before padding. y: A 3-D tensor of shape (1, T, K). mask_interval: a list of tensors of shape (M, 2). contains M mask_start and mask_end. list length is actually 1, because we only support single sample inference for now top_k: (`optional`) int The number of highest probability tokens to keep for top-k-filtering. Default to -100. top_p: (`optional`) float For Neucleus sampling temperature: (`optional`) float The value used to module the next token probabilities. Must be strictly positive. Default to 1.0. eog_coef: (`optional`) float if 0, no change to eog token logits, otherwise, will adjust eog token logit based on the difference between acoustic token and phn token length stop_repetition (`optional`) int if not -1, will set the logits of a token that repeated this many times to be -100000, to avoid generating it again. This only apply to tokens from the first codebook allowed_repeat_tokens (`optional`) list of ints by inspecting the validation set, get a few tokens that indeed repeat a significant amount of time, and exclude those tokens from prevent repetition ultimate_stop_repetition (`optional`) int no matter that token it is, stop repetition once after this number """ assert x.ndim == 2, x.shape assert x_lens.ndim == 1, x_lens.shape assert y.ndim == 3, y.shape if self.args.special_first: y = y + int(self.args.n_special) y = y.transpose(2,1) # [1,T,K] -> [1,K,T] assert y.shape[0] == 1 and y.shape[1] == self.args.n_codebooks, y.shape # there is no padding assert mask_interval.shape == torch.Size((1, mask_interval.shape[1], 2)), mask_interval # make x attention mask and x_input x_attention_mask = torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1).bool().to(x.device) # x_attention_mask = torch.zeros(x.shape[1], x.shape[1]).bool().to(x.device) x_input = self.text_embedding(x) x_input = self.text_positional_embedding(x_input) # make initial y_input # make mask_interval and non_mask_interval y_len = y.shape[2] y_lens = torch.LongTensor([y_len]).to(y.device) mask_interval = mask_interval[0] starts = [item[0].item() for item in mask_interval] + [y_len] ends = [0] + [item[1].item() for item in mask_interval] mask_intervals = [[ (item[0].item(), item[1].item()) for item in mask_interval ]] # a werid name change, mask_interval is input, now is mask_intervals, with one more dimension non_mask_intervals = [[ (ns, ne) for ns, ne in zip(ends, starts) ]] # rearrange y # will add have EOG in each section (SOG will be generated by the pattern class) # but mask can be inserted later after we have shifted the input # y could be rearranged in this way: # [ # [tensor[4, 12], tensor[4, 45], tensor[4, 102], tensor[4, 32]], tensor[4, 22]], # [tensor[4, 44], tensor[4, 56], tensor[4, 19]], # ... # ] # for the first list of tensors (4 tensors), first 3 tensors are non_masked part, last 2 are masked part. # NOTE #non_masked_part = #masked_part + 1 rearranged_y = self.rearrange(y, non_mask_intervals, mask_intervals) assert rearranged_y[0][0].shape[0] == self.args.n_codebooks, rearranged_y[0][0].shape # shift each element of y # next we need to apply pattern shifting to each tensor, after which, we'll replace the starting tokens of each section with a token that's different from the special padding token # [ # [empty, 1, 2, 3, eog, empty, empty, empty], # [empty, empty, 1, 2, 3, eog, empty, empty], # [empty, empty, empty, 1, 2, 3, eog, empty], # [empty, empty, empty, empty, 1, 2, 3, eog] # ] shifted_y, patterns = self.shift(rearranged_y) # each element [K S], patterns is not used, as we directly use the original input y assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape # insert mask token at the intersction of each tensor, but *actually inserted eog as place holder* # the position of inserted mask is also recorded # and the mask_value, the index of the mask emb is recorded inserted_y, mask_position, mask_value = self.insert_mask(shifted_y) assert inserted_y[0][0].shape[0] == self.args.n_codebooks, inserted_y[0][0].shape[0] assert inserted_y[0][1].shape == torch.Size((self.args.n_codebooks, 1)), f"this should be a mask, so should have shape {(self.args.n_codebooks, 1)}, but it's {inserted_y[0][1].shape}" # then concat tensors that belong to the same sample (in order) then get the length of each sample, and then stack them in batch dimension, pad them with pad_token cated_y, new_y_lens = self.cat_y(inserted_y, mask_position, y_lens) # KTB assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], len(inserted_y))) assert not (cated_y == self.args.audio_pad_token).any(), cated_y ### NOTE this is different from forward, as we will remove the masked tokens ### say there are two masked region ### the cated_y should be like ### [empty a a a a mask0 empty b b b mask1 empty c c mask0 empty] ### which means we need to take the part after the last empty out num_mask = len(mask_position[0])//2 assert num_mask == len(mask_position[0])/2, mask_position cated_y = cated_y[:, :mask_position[0][num_mask]+2] # of shape [K,T,B] # logging.info(f"mask_position[0][num_mask]+2: {mask_position[0][num_mask]+2}") more_mask_value = mask_value[0][num_mask+1:] # NOTE this will be used in the generation loop for reference for inserting mask embedding new_y_lens[0] = mask_position[0][num_mask]+2 mask_position[0] = mask_position[0][:num_mask+1] assert mask_position[0][num_mask]+2 == cated_y.shape[1], f"num_mask: {num_mask}, mask_position: {mask_position}, cated_y.shape: {cated_y.shape}" # embed: remember to separately embed the mask tokens embedded_y = self.embed_y(cated_y, mask_position, [mask_value[0][:num_mask+1]]) #BTD # assert embedded_y.shape == torch.Size((y.shape[0], max(new_y_lens), self.args.d_model)), embedded_y.shape # positional embedding y_input = self.audio_positional_embedding(embedded_y) # make attention mask and padding mask y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device) # y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device) x_padding_mask = torch.full((1,x_lens[0]), False).to(x.device) y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device) codebook_eog = [False] * self.args.n_codebooks generated = [] # doesn't contain any empty_token, contains eog cur_generated = [] # say 0 is empty, 4 is eog # tensor([[ 1, 2, 3, 4, 0, 0], # [ 0, 1, 2, 3, 4, 0], # [ 0, 0, 1, 2, 3, 4]]) num_gen = [] cur_num_gen = 0 ##################### silence repetition handling ##################### ##################### silence repetition handling ##################### logging.info(f"silence tokens: {silence_tokens}, note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default") consec_silence_count = 0 prev_token = None ##################### silence repetition handling ##################### ##################### silence repetition handling ##################### # prepare the cache placeholder # n_layers, 2, bsz, num_heads, src_len, head_dim past = torch.ones([self.args.num_decoder_layers, 2, x.shape[0]], device=x.device, dtype=torch.float32) if kvcache else None # handle multi-span kv-cache new_masked_span = False def sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_token, consec_silence_count, stop_repetition, silence_tokens, cur_num_gen): if n_eog == 0: logits_adjust = logits for jj in range(1,self.args.n_codebooks): logits_adjust[jj][self.args.eog] = -10000 logits_adjust[jj][self.args.empty_token] = -10000 ##################### silence repetition handling ##################### if stop_repetition > 0 and prev_token in silence_tokens and consec_silence_count > stop_repetition: if logits_adjust[0, prev_token] < 0: logits_adjust[0, prev_token] = logits_adjust[0, prev_token] * (consec_silence_count - (stop_repetition-1)) else: logits_adjust[0, prev_token] = logits_adjust[0, prev_token] / (consec_silence_count - (stop_repetition-1)) ##################### silence repetition handling ##################### if type(logits_adjust) == list: samples_list= [] for logit in logits_adjust: # print(logit) # print(logit.shape) cur_sample = topk_sampling( logit.unsqueeze(0), top_k=top_k, top_p=top_p, temperature=temperature ) # [1, 1] samples_list.append(cur_sample) samples = torch.cat(samples_list, dim=0) # [K, 1] else: samples = topk_sampling( logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature ) # [K, 1] assert samples.shape == torch.Size((self.args.n_codebooks, 1)), f"samples.shape: {samples.shape}" if cur_num_gen < self.args.n_codebooks-1: for jj in range(1, self.args.n_codebooks - cur_num_gen): samples[-jj, 0] = self.args.empty_token if ( samples[0,0] == self.args.eog or torch.argmax(logits[0], dim=-1) == self.args.eog or y_input.shape[1] > x_lens[0] * 10 ): # last one means y is already too long, shouldn't happen, but put it here samples[0,0] = self.args.eog codebook_eog[0] = True ##################### silence repetition handling ##################### ##################### silence repetition handling ##################### if samples[0,0] in silence_tokens and samples[0,0] == prev_token: consec_silence_count += 1 else: consec_silence_count = 0 prev_token = samples[0,0] ##################### silence repetition handling ##################### ##################### silence repetition handling ##################### return samples, codebook_eog, prev_token, consec_silence_count else: assert sum(codebook_eog[i] for i in range(n_eog)) == n_eog, f"codebook_eog: {codebook_eog}, but n_eog: {n_eog}" logits_adjust = logits for jj in range(n_eog+1,self.args.n_codebooks): logits_adjust[jj][self.args.eog] = -10000 logits_adjust[jj][self.args.empty_token] = -10000 if type(logits_adjust) == list: samples_list= [] for logit in logits_adjust: cur_sample = topk_sampling( logit.unsqueeze(0), top_k=top_k, top_p=top_p, temperature=temperature ) # [1, 1] samples_list.append(cur_sample) samples = torch.cat(samples_list, dim=0) # [K, 1] else: samples = topk_sampling( logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature ) # [K, 1] for jj in range(n_eog): samples[jj, 0] = self.args.empty_token samples[n_eog, 0] = self.args.eog codebook_eog[n_eog] = True return samples, codebook_eog, prev_token, consec_silence_count while True: y_out, present = self.dec_forward( x_input, x_lens, x_attention_mask, x_padding_mask, y_input, new_y_lens, y_attention_mask, y_padding_mask, past=past, last_3_tokens = new_masked_span ) if new_masked_span: new_masked_span = False if past != None: past = torch.cat([past, present.to(past.dtype)], dim=-2) if past.ndim > 3 else present.to(past.dtype) y_out = y_out[:, -1:] # only take the last one logits = torch.stack([self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1) # [B K S card], B==S==1, so [1 K 1 card] logits = logits.squeeze(0).squeeze(1) # [K card] assert logits.shape == torch.Size((self.args.n_codebooks, self.n_audio_tokens[0])), f"{logits.shape}" n_eog = sum(codebook_eog) assert n_eog < self.args.n_codebooks if self.args.eos > 0: # eos stands for end-of-sentence, which shouldn't be used as we are doing speech editing for jj in range(self.args.n_codebooks): logits[jj][self.args.eos] = -10000. # need to use a helper function to hand different n_eog cases samples, codebook_eog, prev_token, consec_silence_count = sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_token, consec_silence_count, stop_repetition, silence_tokens, cur_num_gen) cur_num_gen += 1 cur_generated.append(samples.squeeze(-1)) # [K,1] -> [K] # get samples_emb samples_emb = torch.stack([self.audio_embedding[k](samples[k]) for k in range(self.args.n_codebooks)], dim=0) # [K,1,D] samples_emb = samples_emb.sum(dim=0,keepdim=True) # [1,1,D] if sum(codebook_eog) == self.args.n_codebooks: # generation for the current span is done # re-init codebook_eog = [False] * self.args.n_codebooks num_gen.append(cur_num_gen) cur_num_gen = 0 generated.append(cur_generated) cur_generated = [] # if the current mask span is the last span, then all done # else # append the next mask token and the four empty tokens to start the next generation if len(more_mask_value) > 0: next_mask_ind = more_mask_value.pop(0) mask_emb = self.mask_embedding[next_mask_ind].unsqueeze(0).unsqueeze(0) # [1,1,D] assert mask_emb.shape == torch.Size((1,1,self.args.d_model)), mask_emb.shape empty_token = torch.LongTensor([self.args.empty_token]).to(y.device) empty_emb = torch.stack([ self.audio_embedding[k](empty_token) for k in range(self.args.n_codebooks)], dim=0 ).sum(dim=0, keepdim=True) # [1,1,D] assert empty_emb.shape == torch.Size((1,1,self.args.d_model)), empty_emb.shape extra_emb = torch.cat([mask_emb, empty_emb], dim=1) # [1,2,D] samples_emb = torch.cat([samples_emb, extra_emb], dim=1) # [1,3,D] # prev_last_token, mask_token, empty token assert samples_emb.shape == torch.Size((1,3,self.args.d_model)), f"samples_emb.shape: {samples_emb.shape}" ##################### silence repetition handling ##################### ##################### silence repetition handling ##################### consec_silence_count = 0 prev_token = None ##################### silence repetition handling ##################### ##################### silence repetition handling ##################### # handling kv-caching for multi-span editing new_masked_span = True else: break else: assert samples_emb.shape == torch.Size((1,1,self.args.d_model)), f"samples_emb.shape: {samples_emb.shape}" embedded_y = torch.cat([embedded_y, samples_emb], dim=1) # positional embedding y_input = self.audio_positional_embedding(embedded_y) # [B T D] # make attention mask and padding mask y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device) new_y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device) y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device) assert len(generated) == num_mask, f"len(generated): {len(generated)}, num_mask: {num_mask}" # # combine non_masked_span with generated spans # first need to shift the generated part back flatten_gen = [] for l, orig_span in enumerate(generated): span = torch.stack(orig_span, dim=0) # [T K] span = span.transpose(1,0) # [K, T] assert span.shape[0] == self.args.n_codebooks, span.shape unshifted_span = [] for j, s in enumerate(span): start_from = j end_at = - (self.args.n_codebooks - start_from) unshifted_span.append(s[start_from:end_at]) unshifted_span = torch.stack(unshifted_span, dim=0) assert unshifted_span.shape[1] == num_gen[l] - self.args.n_codebooks, f"len(unshifted_spans[0]): {len(unshifted_span[0])}, num_gen[l]: {num_gen[l]}" flatten_gen.append(unshifted_span) # logging.info(f"unshfited_span: {unshifted_span.shape}") # raise assert len(non_mask_intervals[0]) - 1 == len(flatten_gen), f"len(non_mask_intervals[0]): {len(non_mask_intervals[0])}, len(flatten_gen): {len(flatten_gen)}" res = [] for orig_interval, gen in zip(non_mask_intervals[0], flatten_gen): res.append(y[0, :, orig_interval[0]:orig_interval[1]]) res.append(gen) res.append(y[0, :, non_mask_intervals[0][-1][0]:non_mask_intervals[0][-1][1]]) res = torch.cat(res, dim=1).unsqueeze(0) # [K,new_T] -> [1, K, new_T] expected_y_len = y_len - sum([item[1] - item[0] for item in mask_intervals[0]]) + sum([item - self.args.n_codebooks for item in num_gen]) assert res.shape == torch.Size((1, self.args.n_codebooks, expected_y_len)), f"res.shape: {res.shape}, expected_y_len: {expected_y_len}. y_len - sum([item[1] - item[0] for item in mask_interval]) + sum([item - self.args.n_codebooks for item in num_gen]): {y_len}-{sum([item[1] - item[0] for item in mask_interval])} + {sum([item - self.args.n_codebooks for item in num_gen])}" if self.args.special_first: res = res - int(self.args.n_special) return res def inference_tts( self, x: torch.Tensor, x_lens: torch.Tensor, y: torch.Tensor, top_k: int=-100, top_p: float=1.0, temperature: float=1.0, stop_repetition: int=3, kvcache: int=1, silence_tokens: list[int]=[1388,1898,131], *kargs ) -> torch.Tensor: """ different from inference_tts, this implementation uses kvcache, which should have significant speed up Args: x: A 2-D tensor of shape (1, L). x_lens: A 1-D tensor of shape (1,). It contains the number of tokens in `x` before padding. y: A 3-D tensor of shape (1, T, K). top_k: (`optional`) int The number of highest probability tokens to keep for top-k-filtering. Default to -100. top_p: (`optional`) float For Neucleus sampling temperature: (`optional`) float The value used to module the next token probabilities. Must be strictly positive. Default to 1.0. """ eog_inference = self.args.eos if self.args.eos>0 else self.args.eog assert x.ndim == 2, x.shape assert x_lens.ndim == 1, x_lens.shape assert y.ndim == 3, y.shape if self.args.special_first: y = y + int(self.args.n_special) y = y.transpose(2,1) # [1,T,K] -> [1,K,T] assert y.shape[0] == 1 and y.shape[1] == self.args.n_codebooks, y.shape # there is no padding # make x attention mask and x_input x_attention_mask = torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1).bool().to(x.device) # x_attention_mask = torch.zeros(x.shape[1], x.shape[1]).bool().to(x.device) x_input = self.text_embedding(x) x_input = self.text_positional_embedding(x_input) y_len = y.shape[2] y_lens = torch.LongTensor([y_len]).to(y.device) # rearrange y, we don't add eog to the end, this doesn't actually do anything in the tts scenario rearranged_y = [[y[0]]] assert rearranged_y[0][0].shape[0] == self.args.n_codebooks, rearranged_y[0][0].shape # shift y to create the delayed pattern shifted_y, patterns = self.shift(rearranged_y) # each element [K S], patterns is not used, as we directly use the original input y assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape assert len(shifted_y[0]) == 1, len(shifted_y[0]) # below is different from forward or inference # where we cut this shifted part shifted_y[0][0] = shifted_y[0][0][:, :-(self.args.n_codebooks-1)] assert not (shifted_y[0][0][self.args.n_codebooks:] == self.args.empty_token).any() and not (shifted_y[0][0][self.args.n_codebooks:] == self.args.eog).any(), shifted_y[0][0] # next section in inference is insert mask at the intersection of each tensor in a sample, but we don't need to do that # next section is concate tensors of each sample to one tensor, which we also don't need cated_y = shifted_y[0][0].unsqueeze(-1) #[K,S]->[K,S,B] new_y_lens = torch.LongTensor([cated_y.shape[1]]).to(cated_y.device) assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], 1)) assert not (cated_y == self.args.audio_pad_token).any(), cated_y # replace tokens in y with the embeddings, add sum codebooks up embedded_y = torch.stack([self.audio_embedding[k](cated_y[k]) for k in range(self.args.n_codebooks)], dim=0) # [K, S, B, D] assert embedded_y.shape[0] == self.args.n_codebooks, embedded_y.shape assert embedded_y.shape[-1] == self.args.d_model, embedded_y.shape embedded_y = embedded_y.sum(dim=0) # [K,S,B,D]->[S,B,D] embedded_y = embedded_y.transpose(1,0) # [S,B,D]->[B,S,D] # positional embedding y_input = self.audio_positional_embedding(embedded_y) # make attention mask and padding mask y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device) x_padding_mask = torch.full((1,x_lens[0]), False).to(x.device) y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device) # entering the generation stage # starting from line 708 codebook_eog = [False] * self.args.n_codebooks generated = [] # doesn't contain any empty token, contain eog cur_generated = [] # say 0 is empty, 4 is eog # tensor([[ 1, 2, 3, 4, 0, 0], # [ 0, 1, 2, 3, 4, 0], # [ 0, 0, 1, 2, 3, 4]]) num_gen = [] cur_num_gen = 0 ##################### silence repetition handling ##################### ##################### silence repetition handling ##################### logging.info(f"silence tokens: {silence_tokens}, note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default") consec_silence_count = 0 prev_token = None ##################### silence repetition handling ##################### ##################### silence repetition handling ##################### # prepare the cache placeholder # n_layers, 2, bsz, num_heads, src_len, head_dim past = torch.ones([self.args.num_decoder_layers, 2, x.shape[0]], device=x.device, dtype=torch.float32) if kvcache else None # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}") # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}") # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}") def sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_token, consec_silence_count, stop_repetition, silence_tokens, cur_num_gen): if n_eog == 0: logits_adjust = logits for jj in range(1,self.args.n_codebooks): logits_adjust[jj][eog_inference] = -10000 logits_adjust[jj][self.args.empty_token] = -10000 if cur_num_gen <= self.args.encodec_sr // 5: # this shouldn't happen, but just in case the model stopped too early logits_adjust[0][eog_inference] = -10000 ##################### silence repetition handling ##################### if stop_repetition > 0 and prev_token in silence_tokens and consec_silence_count > stop_repetition: if logits_adjust[0, prev_token] < 0: logits_adjust[0, prev_token] = logits_adjust[0, prev_token] * (consec_silence_count - (stop_repetition-1)) else: logits_adjust[0, prev_token] = logits_adjust[0, prev_token] / (consec_silence_count - (stop_repetition-1)) ##################### silence repetition handling ##################### samples = topk_sampling( logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature ) # [K, 1] assert samples.shape == torch.Size((self.args.n_codebooks, 1)), f"samples.shape: {samples.shape}" if cur_num_gen < self.args.n_codebooks-1: for jj in range(1, self.args.n_codebooks - cur_num_gen): samples[-jj, 0] = self.args.empty_token if ( samples[0,0] == eog_inference or torch.argmax(logits[0], dim=-1) == eog_inference or y_input.shape[1] > x_lens[0] * (self.args.encodec_sr//5) ): # last one means y is already too long, shouldn't happen, but put it here samples[0,0] = eog_inference codebook_eog[0] = True ##################### silence repetition handling ##################### if samples[0,0] in silence_tokens and samples[0,0] == prev_token: consec_silence_count += 1 else: consec_silence_count = 0 prev_token = samples[0,0] ##################### silence repetition handling ##################### return samples, codebook_eog, prev_token, consec_silence_count else: assert sum(codebook_eog[i] for i in range(n_eog)) == n_eog, f"codebook_eog: {codebook_eog}, but n_eog: {n_eog}" logits_adjust = logits for jj in range(n_eog+1,self.args.n_codebooks): logits_adjust[jj][eog_inference] = -10000 logits_adjust[jj][self.args.empty_token] = -10000 samples = topk_sampling( logits_adjust, top_k=top_k, top_p=top_p, temperature=temperature ) # [K, 1] for jj in range(n_eog): samples[jj, 0] = self.args.empty_token samples[n_eog, 0] = eog_inference codebook_eog[n_eog] = True return samples, codebook_eog, prev_token, consec_silence_count while True: y_out, present = self.dec_forward( x_input, x_lens, x_attention_mask, x_padding_mask, y_input, new_y_lens, y_attention_mask, y_padding_mask, past=past ) if past != None: past = torch.cat([past, present.to(past.dtype)], dim=-2) if past.ndim > 3 else present.to(past.dtype) y_out = y_out[:, -1:] # only take the last token logits = torch.stack([self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1) # [B K S card], B==S==1, so [1 K 1 card] logits = logits.squeeze(0).squeeze(1) # [K card] assert logits.shape == torch.Size((self.args.n_codebooks, self.n_audio_tokens[0])), f"{logits.shape}" n_eog = sum(codebook_eog) assert n_eog < self.args.n_codebooks if self.args.eos > 0: # if we are using end-of-sentence token (which is used by default), eog shouldn't be used here, as there is no masked spans for jj in range(self.args.n_codebooks): logits[jj][self.args.eog] = -10000. samples, codebook_eog, prev_token, consec_silence_count = sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_token, consec_silence_count, stop_repetition, silence_tokens, cur_num_gen) cur_num_gen += 1 cur_generated.append(samples.squeeze(-1)) # [K,1] -> [K] # samples.shape is [K,1] # ge samples_emb samples_emb = torch.stack([self.audio_embedding[k](samples[k]) for k in range(self.args.n_codebooks)], dim=0) # [K,1,D] samples_emb = samples_emb.sum(dim=0,keepdim=True) # [1,1,D] if sum(codebook_eog) == self.args.n_codebooks: # generation for the current span is done codebook_eog = [False] * self.args.n_codebooks num_gen.append(cur_num_gen) cur_num_gen = 0 generated.append(cur_generated) cur_generated = [] break else: assert samples_emb.shape == torch.Size((1,1,self.args.d_model)), f"samples_emb.shape: {samples_emb.shape}" embedded_y = torch.cat([embedded_y, samples_emb], dim=1) y_input = self.audio_positional_embedding(embedded_y) # [B T D] # make attention mask and padding mask y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device) new_y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device) y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device) assert len(generated) == 1, f"len(generated): {len(generated)}" # revert the pattern flatten_gen = [] for l, orig_span in enumerate(generated): span = torch.stack(orig_span, dim=0) # [T, K] span = span.transpose(1,0) # [K, T] assert span.shape[0] == self.args.n_codebooks, span.shape unshifted_span = [] for j, s in enumerate(span): start_from = j end_at = - (self.args.n_codebooks - start_from) unshifted_span.append(s[start_from:end_at]) unshifted_span = torch.stack(unshifted_span, dim=0) assert unshifted_span.shape[1] == num_gen[l] - self.args.n_codebooks, f"len(unshifted_spans[0]): {len(unshifted_span[0])}, num_gen[l]: {num_gen[l]}" flatten_gen.append(unshifted_span) assert len(flatten_gen) == 1, len(flatten_gen) # combine res = [y[0], flatten_gen[0]] res = torch.cat(res, dim=1).unsqueeze(0) # [K, new_t] -> [1, K, new_T] expected_y_len = y_len + sum([item - self.args.n_codebooks for item in num_gen]) assert res.shape == torch.Size((1, self.args.n_codebooks, expected_y_len)), f"res.shape: {res.shape}, expected_y_len: {expected_y_len}. y_len + sum([item - self.args.n_codebooks for item in num_gen]): {y_len} + {sum([item - self.args.n_codebooks for item in num_gen])}" if self.args.special_first: res = res - int(self.args.n_special) flatten_gen = flatten_gen - int(self.args.n_special) return res, flatten_gen[0].unsqueeze(0) def inference_tts_batch( self, x: torch.Tensor, x_lens: torch.Tensor, y: torch.Tensor, top_k: int=-100, top_p: float=1.0, temperature: float=1.0, stop_repetition: int=3, kvcache: int=1, batch_size: int=5, silence_tokens: list[int]=[1388,1898,131], *kargs ) -> torch.Tensor: """ have a batch size when forward passing, but they are equivalant to same example but different random seed, therefore as long as one example generated eog, we can drop all other samlpes different from inference_tts, this implementation uses kvcache, which should have significant speed up Args: x: A 2-D tensor of shape (1, L). x_lens: A 1-D tensor of shape (1,). It contains the number of tokens in `x` before padding. y: A 3-D tensor of shape (1, T, K). top_k: (`optional`) int The number of highest probability tokens to keep for top-k-filtering. Default to -100. top_p: (`optional`) float For Neucleus sampling temperature: (`optional`) float The value used to module the next token probabilities. Must be strictly positive. Default to 1.0. """ eog_inference = self.args.eos if self.args.eos>0 else self.args.eog assert x.ndim == 2, x.shape assert x_lens.ndim == 1, x_lens.shape assert y.ndim == 3, y.shape if self.args.special_first: y = y + int(self.args.n_special) y = y.transpose(2,1) # [1,T,K] -> [1,K,T] assert y.shape[0] == 1 and y.shape[1] == self.args.n_codebooks, y.shape # there is no padding # make x attention mask and x_input x_attention_mask = torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1).bool().to(x.device) # x_attention_mask = torch.zeros(x.shape[1], x.shape[1]).bool().to(x.device) x_input = self.text_embedding(x) x_input = self.text_positional_embedding(x_input) y_len = y.shape[2] y_lens = torch.LongTensor([y_len]).to(y.device) # rearrange y, we don't add eog to the end, this doesn't actually do anything in the tts scenario rearranged_y = [[y[0]]] assert rearranged_y[0][0].shape[0] == self.args.n_codebooks, rearranged_y[0][0].shape # shift y to create the delayed pattern shifted_y, patterns = self.shift(rearranged_y) # each element [K S], patterns is not used, as we directly use the original input y assert shifted_y[0][0].shape[0] == self.args.n_codebooks, shifted_y[0][0].shape assert len(shifted_y[0]) == 1, len(shifted_y[0]) # below is different from forward or inference # where we cut this shifted part shifted_y[0][0] = shifted_y[0][0][:, :-(self.args.n_codebooks-1)] assert not (shifted_y[0][0][self.args.n_codebooks:] == self.args.empty_token).any() and not (shifted_y[0][0][self.args.n_codebooks:] == self.args.eog).any(), shifted_y[0][0] # next section in inference is insert mask at the intersection of each tensor in a sample, but we don't need to do that # next section is concate tensors of each sample to one tensor, which we also don't need cated_y = shifted_y[0][0].unsqueeze(-1) #[K,S]->[K,S,B] new_y_lens = torch.LongTensor([cated_y.shape[1]]).to(cated_y.device) assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], 1)) assert not (cated_y == self.args.audio_pad_token).any(), cated_y # replace tokens in y with the embeddings, add sum codebooks up embedded_y = torch.stack([self.audio_embedding[k](cated_y[k]) for k in range(self.args.n_codebooks)], dim=0) # [K, S, B, D] assert embedded_y.shape[0] == self.args.n_codebooks, embedded_y.shape assert embedded_y.shape[-1] == self.args.d_model, embedded_y.shape embedded_y = embedded_y.sum(dim=0) # [K,S,B,D]->[S,B,D] embedded_y = embedded_y.transpose(1,0) # [S,B,D]->[B,S,D] # positional embedding y_input = self.audio_positional_embedding(embedded_y) # make attention mask and padding mask y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device) x_padding_mask = torch.full((1,x_lens[0]), False).to(x.device) y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device) # entering the generation stage # starting from line 708 codebook_eog = [False] * self.args.n_codebooks generated = [] # doesn't contain any empty token, contain eog cur_generated = [[] for _ in range(batch_size)] # say 0 is empty, 4 is eog # tensor([[ 1, 2, 3, 4, 0, 0], # [ 0, 1, 2, 3, 4, 0], # [ 0, 0, 1, 2, 3, 4]]) num_gen = [] cur_num_gen = 0 ##################### silence repetition handling ##################### ##################### silence repetition handling ##################### logging.info(f"silence tokens: {silence_tokens}, note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default") consec_silence_counts = [0 for _ in range(batch_size)] prev_tokens = [None for _ in range(batch_size)] ##################### silence repetition handling ##################### ##################### silence repetition handling ##################### # prepare the cache placeholder # n_layers, 2, bsz, num_heads, src_len, head_dim past = torch.ones([self.args.num_decoder_layers, 2, x.shape[0]], device=x.device, dtype=torch.float32) if kvcache else None # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}") # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}") # logging.info(f"number of decoder layers: {self.args.num_decoder_layers}") keep = None # NOTE: this very important, tells which sample to keep def sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_tokens, consec_silence_counts, stop_repetition, silence_tokens, cur_num_gen, keep): if n_eog == 0: logits_adjust = logits for jj in range(1,self.args.n_codebooks): logits_adjust[:,jj,eog_inference] = -10000 logits_adjust[:,jj,self.args.empty_token] = -10000 if cur_num_gen <= self.args.encodec_sr // 5: # this shouldn't happen, but just in case the model stopped too early logits_adjust[:,:,eog_inference] = -10000 ##################### silence repetition handling ##################### for b in range(batch_size): prev_token = prev_tokens[b] consec_silence_count = consec_silence_counts[b] if stop_repetition > 0 and prev_token in silence_tokens and consec_silence_count > stop_repetition: if logits_adjust[b, 0, prev_token] < 0: logits_adjust[b, 0, prev_token] = logits_adjust[b, 0, prev_token] * (consec_silence_count - (stop_repetition-1)) else: logits_adjust[b, 0, prev_token] = logits_adjust[b, 0, prev_token] / (consec_silence_count - (stop_repetition-1)) ##################### silence repetition handling ##################### samples = topk_sampling( logits_adjust.reshape(batch_size * self.args.n_codebooks, logits_adjust.shape[-1]), top_k=top_k, top_p=top_p, temperature=temperature ) # [B*K, 1] samples = samples.reshape(batch_size, self.args.n_codebooks, 1) assert samples.shape == torch.Size((batch_size, self.args.n_codebooks, 1)), f"samples.shape: {samples.shape}" for b in range(batch_size): if cur_num_gen < self.args.n_codebooks-1: for jj in range(1, self.args.n_codebooks - cur_num_gen): samples[b, -jj, 0] = self.args.empty_token if ( samples[b,0,0] == eog_inference or torch.argmax(logits[b,0], dim=-1) == eog_inference or y_input.shape[1] > x_lens[b] * (self.args.encodec_sr//5) ): # last one means y is already too long, shouldn't happen, but put it here samples[b,0,0] = eog_inference codebook_eog[0] = True keep = b # NOTE keep is a very important variable, we only return this one, note that if eog shows up in two samples, keep will be overwritten by the later one (or the last one) ##################### silence repetition handling ##################### if samples[b,0,0] in silence_tokens and samples[b,0,0] == prev_tokens[b]: consec_silence_counts[b] += 1 else: consec_silence_counts[b] = 0 prev_tokens[b] = samples[b,0,0] ##################### silence repetition handling ##################### return samples, codebook_eog, prev_tokens, consec_silence_counts, keep else: assert sum(codebook_eog[i] for i in range(n_eog)) == n_eog, f"codebook_eog: {codebook_eog}, but n_eog: {n_eog}" logits_adjust = logits for jj in range(n_eog+1,self.args.n_codebooks): logits_adjust[:,jj,eog_inference] = -10000 logits_adjust[:,jj,self.args.empty_token] = -10000 samples = topk_sampling( logits_adjust.reshape(batch_size * self.args.n_codebooks, logits_adjust.shape[-1]), top_k=top_k, top_p=top_p, temperature=temperature ) # [B, K, 1] samples = samples.reshape(batch_size, self.args.n_codebooks, 1) for jj in range(n_eog): samples[keep, jj, 0] = self.args.empty_token samples[keep, n_eog, 0] = eog_inference codebook_eog[n_eog] = True return samples, codebook_eog, prev_tokens, consec_silence_counts, keep while True: # if cur_num_gen > 0, should have everything in kvcache, so only pass in the last token # in the first generation step, we repeat each tensor to make their first dimension of length the batch size if cur_num_gen == 0: assert x_input.ndim == 3 and x_input.shape[0] == 1, x_input.shape assert x_padding_mask.ndim == 2 and x_padding_mask.shape[0] == 1, x_padding_mask.shape assert y_input.ndim == 3 and y_input.shape[0] == 1 and y_input.shape[1] == new_y_lens[0], y_input.shape assert embedded_y.ndim == 3 and embedded_y.shape[0] == 1 and embedded_y.shape[1] == new_y_lens[0], embedded_y.shape x_input = x_input.repeat(batch_size, 1, 1) x_lens = x_lens.repeat(batch_size) # x_attention_mask = x_attention_mask.repeat(batch_size, 1, 1) # no need to work with attention mask, it doesn't contain batch dimension x_padding_mask = x_padding_mask.repeat(batch_size, 1) y_input = y_input.repeat(batch_size, 1, 1) new_y_lens = new_y_lens.repeat(batch_size) # y_attention_mask = y_attention_mask.repeat(batch_size, 1, 1) # no need to work with attention mask, it doesn't contain batch dimension y_padding_mask = y_padding_mask.repeat(batch_size, 1) embedded_y = embedded_y.repeat(batch_size, 1, 1) # will be used to concat with newly generated token embedding past = past.repeat(1, 1, batch_size) if past != None else None else: assert x_input.shape[0] == batch_size and x_padding_mask.shape[0] == batch_size and y_input.shape[0] == batch_size and new_y_lens.shape[0] == batch_size, f"x_input.shape: {x_input.shape}, x_padding_mask.shape: {x_padding_mask.shape}, y_input.shape: {y_input.shape}, new_y_lens.shape: {new_y_lens.shape}" y_out, present = self.dec_forward( x_input, x_lens, x_attention_mask, x_padding_mask, y_input, new_y_lens, y_attention_mask, y_padding_mask, past=past ) if past != None: past = torch.cat([past, present.to(past.dtype)], dim=-2) if past.ndim > 3 else present.to(past.dtype) # if no eog emerges, y_out should have batch size of batch_size if sum(codebook_eog) == 0: assert y_out.shape[0] == batch_size and y_out.ndim == 3, y_out.shape y_out = y_out[:, -1:] # only take the last token logits = torch.stack([self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1) # [B K S card], S==1, so [B K 1 card] logits = logits.squeeze(2) # [B K card] assert logits.shape == torch.Size((batch_size, self.args.n_codebooks, self.n_audio_tokens[0])), f"{logits.shape}" n_eog = sum(codebook_eog) if self.args.eos > 0: for jj in range(self.args.n_codebooks): logits[:,jj,self.args.eog] = -10000. samples, codebook_eog, prev_tokens, consec_silence_counts, keep = sample_helper(n_eog, logits, codebook_eog, top_k, top_p, temperature, prev_tokens, consec_silence_counts, stop_repetition, silence_tokens, cur_num_gen, keep) cur_num_gen += 1 if sum(codebook_eog) == 0: # no eog yet, keep batch_size of samples assert keep == None for b in range(batch_size): cur_generated[b].append(samples[b].squeeze(-1)) elif sum(codebook_eog) == 1: # the first eog just showed up in this step assert keep != None cur_generated = cur_generated[keep] cur_generated.append(samples[keep].squeeze(-1)) else: # we are generating the rest eogs for the 'keep' sample cur_generated.append(samples[keep].squeeze(-1)) # samples.shape is [K,1] # ge samples_emb samples_emb = torch.stack([self.audio_embedding[k](samples[:, k]) for k in range(self.args.n_codebooks)], dim=1) # [B, K,1,D] assert samples_emb.shape == torch.Size([batch_size, self.args.n_codebooks, 1, self.args.d_model]) samples_emb = samples_emb.sum(dim=1,keepdim=False) # [B,1,D] if sum(codebook_eog) == self.args.n_codebooks: # generation for the current span is done codebook_eog = [False] * self.args.n_codebooks num_gen.append(cur_num_gen) cur_num_gen = 0 generated.append(cur_generated) cur_generated = [[] for _ in range(batch_size)] break else: assert samples_emb.shape == torch.Size((batch_size,1,self.args.d_model)), f"samples_emb.shape: {samples_emb.shape}" embedded_y = torch.cat([embedded_y, samples_emb], dim=1) y_input = self.audio_positional_embedding(embedded_y) # [B T D] # make attention mask and padding mask y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device) new_y_lens = torch.LongTensor([y_input.shape[1]]).to(y.device).repeat(batch_size) y_padding_mask = torch.full((batch_size,new_y_lens[0]), False).to(y.device) assert len(generated) == 1, f"len(generated): {len(generated)}" # revert the pattern flatten_gen = [] for l, orig_span in enumerate(generated): span = torch.stack(orig_span, dim=0) # [T, K] span = span.transpose(1,0) # [K, T] assert span.shape[0] == self.args.n_codebooks, span.shape unshifted_span = [] for j, s in enumerate(span): start_from = j end_at = - (self.args.n_codebooks - start_from) unshifted_span.append(s[start_from:end_at]) unshifted_span = torch.stack(unshifted_span, dim=0) assert unshifted_span.shape[1] == num_gen[l] - self.args.n_codebooks, f"len(unshifted_spans[0]): {len(unshifted_span[0])}, num_gen[l]: {num_gen[l]}" flatten_gen.append(unshifted_span) assert len(flatten_gen) == 1, len(flatten_gen) # combine res = [y[0], flatten_gen[0]] res = torch.cat(res, dim=1).unsqueeze(0) # [K, new_t] -> [1, K, new_T] expected_y_len = y_len + sum([item - self.args.n_codebooks for item in num_gen]) assert res.shape == torch.Size((1, self.args.n_codebooks, expected_y_len)), f"res.shape: {res.shape}, expected_y_len: {expected_y_len}. y_len + sum([item - self.args.n_codebooks for item in num_gen]): {y_len} + {sum([item - self.args.n_codebooks for item in num_gen])}" if self.args.special_first: res = res - int(self.args.n_special) flatten_gen = flatten_gen - int(self.args.n_special) return res, flatten_gen[0].unsqueeze(0)