Cloyir/index-tts2-ForDgxSpark
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

优化inference attention mask

Browse Source
This commit is contained in:
yrom 2025-05-17 14:38:01 +08:00
parent cb6c73d391
commit 8f7c1f3e93
1 changed files with 11 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
indextts/gpt/model.py
View File

@ -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:]