From 9b9f7ea61b85611f434771ee75f56d10c3e53fc1 Mon Sep 17 00:00:00 2001 From: James Betker Date: Fri, 17 Dec 2021 10:24:14 -0700 Subject: [PATCH] GptTtsHf: Make the input/target placement easier to reason about --- codes/models/gpt_voice/gpt_tts_hf.py | 40 +++++++++++++--------------- 1 file changed, 18 insertions(+), 22 deletions(-) diff --git a/codes/models/gpt_voice/gpt_tts_hf.py b/codes/models/gpt_voice/gpt_tts_hf.py index 6ac96008..46796fd3 100644 --- a/codes/models/gpt_voice/gpt_tts_hf.py +++ b/codes/models/gpt_voice/gpt_tts_hf.py @@ -52,15 +52,14 @@ class GptTtsHf(nn.Module): self.mel_head = nn.Linear(model_dim, self.NUMBER_MEL_CODES) - def get_logits(self, text_inputs, cond_inputs, mel_targets, get_attns=False): - assert text_inputs.shape[1] <= self.max_symbols_per_phrase - assert cond_inputs.shape[1] <= self.max_conditioning_inputs - assert mel_targets.shape[1] <= self.max_mel_tokens + def build_aligned_inputs_and_targets(self, input, start_token, stop_token): + inp = F.pad(input, (1,0), value=start_token) + tar = F.pad(input, (0,1), value=stop_token) + return inp, tar - - text_targets = F.pad(text_inputs, (1,0), value=self.START_TEXT_TOKEN) - text_emb = self.text_embedding(text_targets) - text_emb = text_emb + self.text_pos_embedding(torch.arange(text_emb.shape[1], device=text_targets.device)) + def get_logits(self, text_inputs, cond_inputs, mel_inputs, get_attns=False): + text_emb = self.text_embedding(text_inputs) + text_emb = text_emb + self.text_pos_embedding(torch.arange(text_emb.shape[1], device=text_inputs.device)) conds = [] for k in range(cond_inputs.shape[1]): @@ -70,9 +69,8 @@ class GptTtsHf(nn.Module): conds = torch.stack(conds, dim=1) conds = conds + self.conditioning_embedding - mel_targets = F.pad(mel_targets, (1,0), value=self.START_MEL_TOKEN) - mel_emb = self.gpt.get_input_embeddings()(mel_targets) - mel_emb = mel_emb + self.mel_pos_embedding(torch.arange(mel_emb.shape[1], device=mel_targets.device)) + mel_emb = self.gpt.get_input_embeddings()(mel_inputs) + mel_emb = mel_emb + self.mel_pos_embedding(torch.arange(mel_emb.shape[1], device=mel_inputs.device)) emb = torch.cat([text_emb, conds, mel_emb], dim=1) gpt_out = self.gpt(inputs_embeds=emb, return_dict=True, output_attentions=get_attns) @@ -80,10 +78,10 @@ class GptTtsHf(nn.Module): return gpt_out.attentions enc = gpt_out.last_hidden_state - text_logits = self.final_norm(enc[:, :self.max_symbols_per_phrase+1]) + text_logits = self.final_norm(enc[:, :text_emb.shape[1]]) text_logits = self.text_head(text_logits) text_logits = text_logits.permute(0,2,1) - mel_logits = self.final_norm(enc[:, -(self.max_mel_tokens+1):]) + mel_logits = self.final_norm(enc[:, -mel_emb.shape[1]:]) mel_logits = self.mel_head(mel_logits) mel_logits = mel_logits.permute(0,2,1) @@ -103,13 +101,12 @@ class GptTtsHf(nn.Module): if mel_lengths[b] < mel_targets.shape[-1]: mel_targets[b, mel_lengths[b]:] = self.STOP_MEL_TOKEN - text_logits, mel_logits = self.get_logits(text_inputs, cond_inputs, mel_targets, get_attns=return_attentions) + text_inputs, text_targets = self.build_aligned_inputs_and_targets(text_inputs, self.START_TEXT_TOKEN, self.STOP_TEXT_TOKEN) + mel_inputs, mel_targets = self.build_aligned_inputs_and_targets(mel_targets, self.START_MEL_TOKEN, self.STOP_MEL_TOKEN) + text_logits, mel_logits = self.get_logits(text_inputs, cond_inputs, mel_inputs, get_attns=return_attentions) if return_attentions: return mel_logits - - text_targets = F.pad(text_inputs, (0,self.max_symbols_per_phrase-text_inputs.shape[1]+1), value=self.STOP_TEXT_TOKEN) loss_text = F.cross_entropy(text_logits, text_targets.long()) - mel_targets = F.pad(mel_targets, (0,self.max_mel_tokens-mel_targets.shape[1]+1), value=self.STOP_MEL_TOKEN) loss_mel = F.cross_entropy(mel_logits, mel_targets.long()) return loss_text.mean(), loss_mel.mean(), mel_logits @@ -117,10 +114,9 @@ class GptTtsHf(nn.Module): if not hasattr(self, 'inference_model'): self.inference_model = GPT2InferenceModel(self.gpt_config, self.gpt, self.text_pos_embedding, self.final_norm, self.text_head) - text_targets = F.pad(text_inputs, (1,0), value=self.START_TEXT_TOKEN) - text_targets = F.pad(text_targets, (0,1), value=self.STOP_TEXT_TOKEN) - text_emb = self.text_embedding(text_targets) - text_emb = text_emb + self.text_pos_embedding(torch.arange(text_emb.shape[1], device=text_targets.device)) + text_inputs, text_targets = self.build_aligned_inputs_and_targets(text_inputs, self.START_TEXT_TOKEN, self.STOP_TEXT_TOKEN) + text_emb = self.text_embedding(text_inputs) + text_emb = text_emb + self.text_pos_embedding(torch.arange(text_emb.shape[1], device=text_inputs.device)) conds = [] for k in range(cond_inputs.shape[1]): @@ -133,7 +129,7 @@ class GptTtsHf(nn.Module): emb = torch.cat([text_emb, conds], dim=1) self.inference_model.store_mel_emb(emb) - fake_inputs = torch.full((text_inputs.shape[0],emb.shape[1]+1,), fill_value=1, dtype=torch.long, device=text_inputs.device) + fake_inputs = torch.full((emb.shape[0],emb.shape[1]+1,), fill_value=1, dtype=torch.long, device=text_inputs.device) fake_inputs[:,-1] = self.START_MEL_TOKEN gen = self.inference_model.generate(fake_inputs, do_sample=do_sample, bos_token_id=self.START_MEL_TOKEN, pad_token_id=self.STOP_MEL_TOKEN, eos_token_id=self.STOP_MEL_TOKEN,