From 8f7c1f3e93f485b5bc4b90c245643aff7cf981c8 Mon Sep 17 00:00:00 2001 From: yrom Date: Sat, 17 May 2025 14:38:01 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96inference=20attention=20mask?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- indextts/gpt/model.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/indextts/gpt/model.py b/indextts/gpt/model.py index 2906b31..5852650 100644 --- a/indextts/gpt/model.py +++ b/indextts/gpt/model.py @@ -592,10 +592,11 @@ class UnifiedVoice(nn.Module): def inference_speech(self, speech_conditioning_latent, text_inputs, cond_mel_lengths=None, input_tokens=None, num_return_sequences=1, max_generate_length=None, typical_sampling=False, typical_mass=.9, **hf_generate_kwargs): + text_masks = ((text_inputs != self.stop_text_token) & (text_inputs != self.start_text_token)).long() text_inputs = F.pad(text_inputs, (0, 1), value=self.stop_text_token) text_inputs, _ = self.build_aligned_inputs_and_targets(text_inputs, self.start_text_token, self.stop_text_token) + text_masks = F.pad(text_masks, (1, 1), value=1) # (-1, +1) text_emb = self.text_embedding(text_inputs) + self.text_pos_embedding(text_inputs) - speech_conditioning_latent = self.get_conditioning(speech_conditioning_latent, cond_mel_lengths) conds = speech_conditioning_latent emb = torch.cat([conds, text_emb], dim=1) @@ -604,7 +605,11 @@ class UnifiedVoice(nn.Module): # +1 for the start_audio_token fake_inputs = torch.full((emb.shape[0], emb.shape[1] + 1,), fill_value=1, dtype=torch.long, device=text_inputs.device) - + attention_mask = torch.cat([ + torch.ones((conds.shape[0], conds.shape[1]), dtype=torch.long, device=text_inputs.device), + text_masks, + torch.ones((conds.shape[0], 1), dtype=torch.long, device=text_inputs.device), + ], dim=1) fake_inputs[:, -1] = self.start_mel_token trunc_index = fake_inputs.shape[1] if input_tokens is None: @@ -614,12 +619,15 @@ class UnifiedVoice(nn.Module): 0] == 0, "The number of return sequences must be divisible by the number of input sequences" fake_inputs = fake_inputs.repeat(num_return_sequences, 1) input_tokens = input_tokens.repeat(num_return_sequences // input_tokens.shape[0], 1) + input_tokens_mask = ((input_tokens != self.stop_text_token) & (input_tokens != self.start_text_token)).long() inputs = torch.cat([fake_inputs, input_tokens], dim=1) + attention_mask = torch.cat([attention_mask.repeat(num_return_sequences, 1), input_tokens_mask], dim=1) + logits_processor = LogitsProcessorList([TypicalLogitsWarper(mass=typical_mass)]) if typical_sampling else LogitsProcessorList() max_length = trunc_index + self.max_mel_tokens - 1 if max_generate_length is None else trunc_index + max_generate_length gen = self.inference_model.generate(inputs, bos_token_id=self.start_mel_token, pad_token_id=self.stop_mel_token, - eos_token_id=self.stop_mel_token, + eos_token_id=self.stop_mel_token, attention_mask=attention_mask, max_length=max_length, logits_processor=logits_processor, num_return_sequences=num_return_sequences, **hf_generate_kwargs) return gen[:, trunc_index:]