From c68669e1e1e8b57e3eeb324c5f29a6d1bb3b3b71 Mon Sep 17 00:00:00 2001 From: James Betker Date: Tue, 14 Jun 2022 15:18:58 -0600 Subject: [PATCH] uv2 add alignment head --- codes/models/audio/tts/unified_voice2.py | 84 ++++++++++++------------ 1 file changed, 42 insertions(+), 42 deletions(-) diff --git a/codes/models/audio/tts/unified_voice2.py b/codes/models/audio/tts/unified_voice2.py index 134147cd..fa85ffc6 100644 --- a/codes/models/audio/tts/unified_voice2.py +++ b/codes/models/audio/tts/unified_voice2.py @@ -238,22 +238,7 @@ class MelEncoder(nn.Module): class UnifiedVoice(nn.Module): def __init__(self, layers=8, model_dim=512, heads=8, max_text_tokens=120, max_mel_tokens=250, max_conditioning_inputs=1, mel_length_compression=1024, number_text_tokens=256, number_mel_codes=8194, start_mel_token=8192, - stop_mel_token=8193, start_text_token=None, checkpointing=True, types=1): - """ - Args: - layers: Number of layers in transformer stack. - model_dim: Operating dimensions of the transformer - heads: Number of transformer heads. Must be divisible by model_dim. Recommend model_dim//64 - max_text_tokens: Maximum number of text tokens that will be encountered by model. - max_mel_tokens: Maximum number of MEL tokens that will be encountered by model. - max_conditioning_inputs: Maximum number of conditioning inputs provided to the model. If (1), conditioning input can be of format (b,80,s), otherwise (b,n,80,s). - mel_length_compression: The factor between and . Used to compute MEL code padding given wav input length. - number_text_tokens: - number_mel_codes: - start_mel_token: - stop_mel_token: - checkpointing: - """ + stop_mel_token=8193, start_text_token=255, checkpointing=True, types=1, only_alignment_head=False): super().__init__() self.number_text_tokens = number_text_tokens @@ -278,6 +263,15 @@ class UnifiedVoice(nn.Module): self.final_norm = nn.LayerNorm(model_dim) self.text_head = nn.Linear(model_dim, self.number_text_tokens*types+1) self.mel_head = nn.Linear(model_dim, self.number_mel_codes) + self.alignment_head = nn.Linear(model_dim, 256) + + if only_alignment_head: + for p in self.parameters(): + p.DO_NOT_TRAIN = True + p.requires_grad = False + for p in self.alignment_head.parameters(): + del p.DO_NOT_TRAIN + p.requires_grad = True # Initialize the embeddings per the GPT-2 scheme embeddings = [self.text_embedding, self.mel_embedding] @@ -310,11 +304,8 @@ class UnifiedVoice(nn.Module): mel_input_tokens[b, actual_end:] = self.stop_mel_token return mel_input_tokens - def get_logits(self, speech_conditioning_inputs, first_inputs, first_head, second_inputs=None, second_head=None, return_latent=False): - if second_inputs is not None: - emb = torch.cat([speech_conditioning_inputs, first_inputs, second_inputs], dim=1) - else: - emb = torch.cat([speech_conditioning_inputs, first_inputs], dim=1) + def get_logits(self, speech_conditioning_inputs, first_inputs, second_inputs, return_latent=False): + emb = torch.cat([speech_conditioning_inputs, first_inputs, second_inputs], dim=1) gpt_out = self.gpt(inputs_embeds=emb, return_dict=True) @@ -324,16 +315,19 @@ class UnifiedVoice(nn.Module): if return_latent: return enc[:, speech_conditioning_inputs.shape[1]:speech_conditioning_inputs.shape[1]+first_inputs.shape[1]], enc[:, -second_inputs.shape[1]:] - first_logits = enc[:, :first_inputs.shape[1]] - first_logits = first_head(first_logits) - first_logits = first_logits.permute(0,2,1) - if second_inputs is not None: - second_logits = enc[:, -second_inputs.shape[1]:] - second_logits = second_head(second_logits) - second_logits = second_logits.permute(0,2,1) - return first_logits, second_logits - else: - return first_logits + text_logits = enc[:, :first_inputs.shape[1]] + text_logits = self.text_head(text_logits) + text_logits = text_logits.permute(0,2,1) + + mel_logits = enc[:, -second_inputs.shape[1]:] + mel_logits = self.mel_head(mel_logits) + mel_logits = mel_logits.permute(0,2,1) + + alignment_logits = enc[:, -second_inputs.shape[1]:] + alignment_logits = self.alignment_head(alignment_logits) + alignment_logits = alignment_logits.permute(0,2,1) + + return text_logits, mel_logits, alignment_logits def get_conditioning_latent(self, speech_conditioning_input): @@ -346,7 +340,7 @@ class UnifiedVoice(nn.Module): return conds - def forward(self, speech_conditioning_input, text_inputs, text_lengths, mel_codes, wav_lengths, types=None, text_first=True, return_latent=False): + def forward(self, speech_conditioning_input, text_inputs, text_lengths, mel_codes, ctc_codes, wav_lengths, types=None, return_latent=False): """ Forward pass that uses both text and voice in either text conditioning mode or voice conditioning mode (actuated by `text_first`). @@ -363,12 +357,22 @@ class UnifiedVoice(nn.Module): if types is not None: text_inputs = text_inputs * (1+types).unsqueeze(-1) + # TODO: do this in the dataloader. + for b in range(ctc_codes.shape[0]): + last_code = 0 + for j in range(ctc_codes.shape[1]): + if ctc_codes[b][j] == 0: + ctc_codes[b][j] = last_code + else: + last_code = ctc_codes[b][j] + alignment_targets = F.interpolate(ctc_codes.unsqueeze(1).float(), size=(mel_codes.shape[-1],), mode='nearest').long().squeeze() + mel_codes = self.set_mel_padding(mel_codes, wav_lengths) text_inputs = F.pad(text_inputs, (0,1), value=self.stop_text_token) mel_codes = F.pad(mel_codes, (0,1), value=self.stop_mel_token) + alignment_targets = F.pad(alignment_targets, (0,2), value=0) conds = self.get_conditioning_latent(speech_conditioning_input) - 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) + self.text_pos_embedding(text_inputs) mel_codes, mel_targets = self.build_aligned_inputs_and_targets(mel_codes, self.start_mel_token, self.stop_mel_token) @@ -376,18 +380,14 @@ class UnifiedVoice(nn.Module): mel_emb = self.mel_embedding(mel_inp) mel_emb = mel_emb + self.mel_pos_embedding(mel_codes) - if text_first: - text_logits, mel_logits = self.get_logits(conds, text_emb, self.text_head, mel_emb, self.mel_head, return_latent=return_latent) - if return_latent: - return mel_logits[:, :-2] # Despite the name, these are not logits. Strip off the two tokens added by this forward pass. - else: - mel_logits, text_logits = self.get_logits(conds, mel_emb, self.mel_head, text_emb, self.text_head, return_latent=return_latent) - if return_latent: - return text_logits[:, :-2] # Despite the name, these are not logits. Strip off the two tokens added by this forward pass. + text_logits, mel_logits, alignment_logits = self.get_logits(conds, text_emb, mel_emb, return_latent=return_latent) + if return_latent: + return mel_logits[:, :-2] # Despite the name, these are not logits. Strip off the two tokens added by this forward pass. loss_text = F.cross_entropy(text_logits, text_targets.long()) loss_mel = F.cross_entropy(mel_logits, mel_targets.long()) - return loss_text.mean(), loss_mel.mean(), mel_logits + loss_alignment = F.cross_entropy(alignment_logits, alignment_targets) + return loss_text.mean(), loss_mel.mean(), loss_alignment, mel_logits def inference_speech(self, speech_conditioning_input, text_inputs, **hf_generate_kwargs): if self.max_mel_tokens == -1: # Assume if this is the case, max_mel_tokens=-1 also