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

修正attention mask和positional embeddings

Browse Source 
- 将之前只有text右侧填充改为cond+text 整体左侧填充
- 添加填充测试用例
This commit is contained in:
yrom 2025-05-18 15:27:53 +08:00
parent a50cb8c287
commit 22eeb7625f
3 changed files with 230 additions and 108 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

176
indextts/gpt/model.py
View File

@ -40,6 +40,7 @@ class ResBlock(nn.Module):
class GPT2InferenceModel(GPT2PreTrainedModel):
def __init__(self, config, gpt, text_pos_emb, embeddings, norm, linear, kv_cache=False):
super().__init__(config)
# Note: the argument named `text_pos_emb` here actually represents the mel position embedding
self.transformer = gpt
self.text_pos_embedding = text_pos_emb
self.embeddings = embeddings
@ -97,7 +98,7 @@ class GPT2InferenceModel(GPT2PreTrainedModel):
if attention_mask is not None and position_ids is None:
# create position_ids on the fly for batch generation
position_ids = attention_mask.long().cumsum(-1) - 1
position_ids.masked_fill_(attention_mask == 0, 1)
position_ids.masked_fill_(attention_mask == 0, 0)
if past_key_values:
position_ids = position_ids[:, -1].unsqueeze(-1)
else:
@ -134,7 +135,6 @@ class GPT2InferenceModel(GPT2PreTrainedModel):
return_dict = (
return_dict if return_dict is not None else self.config.use_return_dict
)
# Create embedding
mel_len = self.cached_mel_emb.shape[1]
if input_ids.shape[1] != 1:
@ -587,80 +587,116 @@ class UnifiedVoice(nn.Module):
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
@staticmethod
def pad_text_masks(text_masks):
def prepare_gpt_inputs(
self,
conditional_latents: torch.Tensor,
text_inputs: torch.Tensor,
):
"""
pad masks to [b, t+2]. e.g.:
```
[
[1, 1, 1, 1, 1] -> [1, 1, 1, 1, 1, 1, 1]
[1, 1, 1, 1, 0] -> [1, 1, 1, 1, 1, 1, 0]
[0, 0, 0, 0, 0] -> [0, 0, 0, 0, 0, 1, 1]
[1, 0, 0, 0, 0] -> [1, 1, 1, 0, 0, 0, 0]
[0, 0, 1, 1, 1] -> [0, 0, 1, 1, 1, 1, 1]
Prepare the inputs for the GPT2InferenceModel to generate.
Args:
conds_latent: (b, 32, dim) audio conditioning embedding by `get_conditioning()`
text_inputs: (b, L)
Returns:
input_ids: (b, s+1) the input ids for the GPT2InferenceModel.generate()
inputs_embeds: (b, s+1, dim) the input embeddings for the GPT2InferenceModel.forward()
attention_mask: (b, s+1) the attention mask for the GPT2InferenceModel.generate()
"""
b, L = text_inputs.shape[:2]
device = text_inputs.device
single_cond = conditional_latents.ndim == 3 and conditional_latents.shape[0] == 1
if not single_cond:
assert conditional_latents.shape[0] == b, f"batch size mismatch: {conditional_latents.shape[0]} vs {b}"
batched_mel_emb = []
attention_masks = []
target_len = conditional_latents.shape[1] + L + 2
for i in range(b):
valid_mask = (text_inputs[i] != self.stop_text_token) & (text_inputs[i] != self.start_text_token)
text_input = text_inputs[i][valid_mask]
text_input = F.pad(text_input, (1, 0), value=self.start_text_token)
text_input = F.pad(text_input, (0, 1), value=self.stop_text_token)
text_input_pos = torch.arange(0, text_input.size(-1), device=device)
text_emb = self.text_embedding(text_input) + self.text_pos_embedding.emb(text_input_pos)
# concatenate [conditional latents][text embeddings]
conds_text_emb = [
conditional_latents.squeeze(0) if single_cond else conditional_latents[i],
text_emb,
]
```
text_masks: [b, t]
padded_masks: [b, t+2]
"""
b, t = text_masks.size()
padded = torch.nn.functional.pad(text_masks, (1, 1), value=0) # [b, t+2]
padded = torch.ones_like(padded, dtype=text_masks.dtype, device=text_masks.device)
# Find the first and last non-zero index
# and set the values before the first index and after the last index to 0
nonzero_mask = text_masks != 0
has_nonzero = nonzero_mask.any(dim=1)
first_idx = nonzero_mask.float().argmax(dim=1)
rev_mask = nonzero_mask.flip(dims=[1])
last_idx = t - rev_mask.float().argmax(dim=1)
first_idx = first_idx.reshape(b, 1)
last_idx = last_idx.reshape(b, 1)
col_idx = torch.arange(t+2, device=text_masks.device).unsqueeze(0).expand(b, -1) # [b, t_2]
padded[col_idx < first_idx] = 0
padded[col_idx > last_idx+1] = 0
# all zeros
padded[~has_nonzero, :-2] = 0
return padded
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 = UnifiedVoice.pad_text_masks(text_masks) # [b, t+2]
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)
self.inference_model.store_mel_emb(emb)
# +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)
# +1 for the start_mel_token
attention_mask = torch.ones(target_len+1, dtype=torch.long, device=device)
# check this text input is padded
padding: int = L + 2 - text_input.size(-1)
# pad left of [cond][text] -> [pad][cond][text]
if padding > 0:
pad = torch.zeros((padding, conditional_latents.size(-1)), dtype=text_emb.dtype, device=device) # [p, dim]
conds_text_emb.insert(0, pad)
attention_mask[:padding] = 0
mel_emb = torch.cat(conds_text_emb) #[s, dim]
assert mel_emb.shape[0] == target_len, f"mel_emb.shape: {mel_emb.shape}, target_len: {target_len}"
batched_mel_emb.append(mel_emb)
attention_masks.append(attention_mask)
# [b, s, dim]
batched_mel_emb = torch.stack(batched_mel_emb, dim=0)
# [b, s+1]
attention_mask = torch.stack(attention_masks, dim=0)
# [b, s+1]
fake_inputs = torch.ones(
(
batched_mel_emb.shape[0],
batched_mel_emb.shape[1] + 1, # +1 for the start_mel_token
),
dtype=torch.long,
device=device,
)
fake_inputs[:, -1] = self.start_mel_token
trunc_index = fake_inputs.shape[1]
return fake_inputs, batched_mel_emb, attention_mask
def inference_speech(self, speech_conditioning_mel, 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):
"""
Args:
speech_conditioning_mel: (b, n_mels, frames) or (n_mels, frames)
text_inputs: (b, L)
cond_mel_lengths: lengths of the conditioning mel spectrograms in shape (b,) or (1,)
input_tokens: additional tokens for generation in shape (b, s) or (s,)
max_generate_length: limit the number of generated tokens
hf_generate_kwargs: kwargs for `GPT2InferenceModel.generate(**hf_generate_kwargs)`
"""
if speech_conditioning_mel.ndim == 2:
speech_conditioning_mel = speech_conditioning_mel.unsqueeze(0)
if cond_mel_lengths is None:
cond_mel_lengths = torch.tensor([speech_conditioning_mel.shape[-1]], device=speech_conditioning_mel.device)
conds_latent = self.get_conditioning(speech_conditioning_mel, cond_mel_lengths)
input_ids, inputs_embeds, attention_mask = self.prepare_gpt_inputs(conds_latent, text_inputs)
self.inference_model.store_mel_emb(inputs_embeds)
if input_tokens is None:
inputs = fake_inputs
inputs = input_ids
else:
assert num_return_sequences % input_tokens.shape[
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)
if input_tokens.ndim == 1:
input_tokens = input_tokens.unsqueeze(0)
assert num_return_sequences % input_tokens.shape[0] == 0, \
"The num_return_sequences must be divisible by the batch number of input_tokens"
assert num_return_sequences % text_inputs.shape[0] == 0, \
"The num_return_sequences must be divisible by the batch number of text_inputs"
b = num_return_sequences // input_ids.shape[0]
if b > 1:
input_ids = input_ids.repeat(b, 1)
attention_mask = attention_mask.repeat(b, 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)
inputs = torch.cat([input_ids, input_tokens], dim=1)
attention_mask = F.pad(attention_mask, (0, input_tokens.shape[1]), value=1)
trunc_index = inputs.shape[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,
max_length = (trunc_index + self.max_mel_tokens - 1) if max_generate_length is None else trunc_index + max_generate_length
output = 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, 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:]
num_return_sequences=num_return_sequences,
**hf_generate_kwargs)
if isinstance(output, torch.Tensor):
return output[:, trunc_index:]
# GenerateOutput
output.sequences = output.sequences[:, trunc_index:]
return output

74
indextts/infer.py
View File

@ -242,7 +242,7 @@ class IndexTTS:
def pad_tokens_cat(self, tokens: List[torch.Tensor]) -> torch.Tensor:
if self.model_version and self.model_version >= 1.5:
# 1.5版本以上，直接使用stop_text_token 右侧填充
# 1.5版本以上，直接使用stop_text_token 右侧填充，填充到最大长度
# [1, N] -> [N,]
tokens = [t.squeeze(0) for t in tokens]
return pad_sequence(tokens, batch_first=True, padding_value=self.cfg.gpt.stop_text_token, padding_side="right")
@ -273,7 +273,7 @@ class IndexTTS:
self.gr_progress(value, desc=desc)
# 快速推理：对于“多句长文本”，可实现至少 2~10 倍以上的速度提升~ （First modified by sunnyboxs 2025-04-16）
def infer_fast(self, audio_prompt, text, output_path, verbose=False, max_text_tokens_per_sentence=100, sentences_bucket_max_size=4):
def infer_fast(self, audio_prompt, text, output_path, verbose=False, max_text_tokens_per_sentence=100, sentences_bucket_max_size=4, **sample_kwargs):
"""
Args:
``max_text_tokens_per_sentence``: 分句的最大token数，默认``100``，可以根据GPU硬件情况调整
@ -321,15 +321,15 @@ class IndexTTS:
print(" splited sentences count:", len(sentences))
print(" max_text_tokens_per_sentence:", max_text_tokens_per_sentence)
print(*sentences, sep="\n")
top_p = 0.8
top_k = 30
temperature = 1.0
do_sample = sample_kwargs.pop("do_sample", True)
top_p = sample_kwargs.pop("top_p", 0.8)
top_k = sample_kwargs.pop("top_k", 30)
temperature = sample_kwargs.pop("temperature", 1.0)
autoregressive_batch_size = 1
length_penalty = 0.0
num_beams = 3
repetition_penalty = 10.0
max_mel_tokens = 600
length_penalty = sample_kwargs.pop("length_penalty", 0.0)
num_beams = sample_kwargs.pop("num_beams", 3)
repetition_penalty = sample_kwargs.pop("repetition_penalty", 10.0)
max_mel_tokens = sample_kwargs.pop("max_mel_tokens", 600)
sampling_rate = 24000
# lang = "EN"
# lang = "ZH"
@ -365,35 +365,31 @@ class IndexTTS:
# Sequential processing of bucketing data
all_batch_num = 0
all_batch_num = sum(len(s) for s in all_sentences)
all_batch_codes = []
processed_num = 0
for item_tokens in all_text_tokens:
batch_num = len(item_tokens)
if batch_num > 1:
batch_text_tokens = self.pad_tokens_cat(item_tokens)
batch_cond_mel_lengths = cond_mel_lengths.expand(batch_num) # [batch_num]
batch_auto_conditioning = auto_conditioning.expand(batch_num, -1, -1) # [batch_num, n_mels, L]
else:
batch_text_tokens = item_tokens[0]
batch_cond_mel_lengths = cond_mel_lengths
batch_auto_conditioning = auto_conditioning
all_batch_num += batch_num
batch_text_tokens = torch.nn.functional.pad(item_tokens[0], (8, 0), value=self.cfg.gpt.start_text_token)
processed_num += batch_num
# gpt speech
self._set_gr_progress(0.2, "gpt inference speech...")
self._set_gr_progress(0.2 + 0.3 * processed_num/all_batch_num, f"gpt inference speech... {processed_num}/{all_batch_num}")
m_start_time = time.perf_counter()
with torch.no_grad():
with torch.amp.autocast(batch_text_tokens.device.type, enabled=self.dtype is not None, dtype=self.dtype):
temp_codes = self.gpt.inference_speech(batch_auto_conditioning, batch_text_tokens,
cond_mel_lengths=batch_cond_mel_lengths,
temp_codes = self.gpt.inference_speech(auto_conditioning, batch_text_tokens,
cond_mel_lengths=cond_mel_lengths,
# text_lengths=text_len,
do_sample=True,
do_sample=do_sample,
top_p=top_p,
top_k=top_k,
temperature=temperature,
num_return_sequences=autoregressive_batch_size if batch_num == 1 else 1,
num_return_sequences=autoregressive_batch_size,
length_penalty=length_penalty,
num_beams=num_beams if batch_num == 1 else 1,
num_beams=num_beams,
repetition_penalty=repetition_penalty,
max_generate_length=max_mel_tokens)
all_batch_codes.append(temp_codes)
@ -490,7 +486,7 @@ class IndexTTS:
return (sampling_rate, wav_data)
# 原始推理模式
def infer(self, audio_prompt, text, output_path, verbose=False, max_text_tokens_per_sentence=120):
def infer(self, audio_prompt, text, output_path, verbose=False, max_text_tokens_per_sentence=120, **sample_kwargs):
print(">> start inference...")
self._set_gr_progress(0, "start inference...")
if verbose:
@ -516,6 +512,7 @@ class IndexTTS:
cond_mel_frame = cond_mel.shape[-1]
pass
self._set_gr_progress(0.1, "text processing...")
auto_conditioning = cond_mel
text_tokens_list = self.tokenizer.tokenize(text)
sentences = self.tokenizer.split_sentences(text_tokens_list, max_text_tokens_per_sentence)
@ -524,14 +521,15 @@ class IndexTTS:
print("sentences count:", len(sentences))
print("max_text_tokens_per_sentence:", max_text_tokens_per_sentence)
print(*sentences, sep="\n")
top_p = 0.8
top_k = 30
temperature = 1.0
do_sample = sample_kwargs.pop("do_sample", True)
top_p = sample_kwargs.pop("top_p", 0.8)
top_k = sample_kwargs.pop("top_k", 30)
temperature = sample_kwargs.pop("temperature", 1.0)
autoregressive_batch_size = 1
length_penalty = 0.0
num_beams = 3
repetition_penalty = 10.0
max_mel_tokens = 600
length_penalty = sample_kwargs.pop("length_penalty", 0.0)
num_beams = sample_kwargs.pop("num_beams", 3)
repetition_penalty = sample_kwargs.pop("repetition_penalty", 10.0)
max_mel_tokens = sample_kwargs.pop("max_mel_tokens", 600)
sampling_rate = 24000
# lang = "EN"
# lang = "ZH"
@ -539,7 +537,7 @@ class IndexTTS:
gpt_gen_time = 0
gpt_forward_time = 0
bigvgan_time = 0
progress = 0
for sent in sentences:
text_tokens = self.tokenizer.convert_tokens_to_ids(sent)
text_tokens = torch.tensor(text_tokens, dtype=torch.int32, device=self.device).unsqueeze(0)
@ -555,7 +553,8 @@ class IndexTTS:
# text_len = torch.IntTensor([text_tokens.size(1)], device=text_tokens.device)
# print(text_len)
progress += 1
self._set_gr_progress(0.2 + 0.4 * (progress-1) / len(sentences), f"gpt inference latent... {progress}/{len(sentences)}")
m_start_time = time.perf_counter()
with torch.no_grad():
with torch.amp.autocast(text_tokens.device.type, enabled=self.dtype is not None, dtype=self.dtype):
@ -563,7 +562,7 @@ class IndexTTS:
cond_mel_lengths=torch.tensor([auto_conditioning.shape[-1]],
device=text_tokens.device),
# text_lengths=text_len,
do_sample=True,
do_sample=do_sample,
top_p=top_p,
top_k=top_k,
temperature=temperature,
@ -587,7 +586,7 @@ class IndexTTS:
print(codes, type(codes))
print(f"fix codes shape: {codes.shape}, codes type: {codes.dtype}")
print(f"code len: {code_lens}")
self._set_gr_progress(0.2 + 0.4 * progress / len(sentences), f"gpt inference speech... {progress}/{len(sentences)}")
m_start_time = time.perf_counter()
# latent, text_lens_out, code_lens_out = \
with torch.amp.autocast(text_tokens.device.type, enabled=self.dtype is not None, dtype=self.dtype):
@ -605,11 +604,12 @@ class IndexTTS:
wav = wav.squeeze(1)
wav = torch.clamp(32767 * wav, -32767.0, 32767.0)
if verbose:
print(f"wav shape: {wav.shape}", "min:", wav.min(), "max:", wav.max())
# wavs.append(wav[:, :-512])
wavs.append(wav.cpu()) # to cpu before saving
end_time = time.perf_counter()
self._set_gr_progress(0.9, "save audio...")
wav = torch.cat(wavs, dim=1)
wav_length = wav.shape[-1] / sampling_rate
print(f">> Reference audio length: {cond_mel_frame * 256 / sampling_rate:.2f} seconds")

86
tests/padding_test.py Normal file
View File

@ -0,0 +1,86 @@
import torch
import torchaudio
from indextts.infer import IndexTTS
from indextts.utils.feature_extractors import MelSpectrogramFeatures
from torch.nn import functional as F
if __name__ == "__main__":
import transformers
transformers.set_seed(42)
audio_prompt="tests/sample_prompt.wav"
tts = IndexTTS(cfg_path="checkpoints/config.yaml", model_dir="checkpoints", is_fp16=False, use_cuda_kernel=False)
text = "晕 XUAN4 是 一 种 not very good GAN3 觉"
text_tokens = tts.tokenizer.encode(text)
text_tokens = torch.tensor(text_tokens, dtype=torch.int32, device=tts.device).unsqueeze(0) # [1, L]
audio, sr = torchaudio.load(audio_prompt)
audio = torch.mean(audio, dim=0, keepdim=True)
audio = torchaudio.transforms.Resample(sr, 24000)(audio)
auto_conditioning = MelSpectrogramFeatures()(audio).to(tts.device)
cond_mel_lengths = torch.tensor([auto_conditioning.shape[-1]]).to(tts.device)
with torch.no_grad():
kwargs = {
"cond_mel_lengths": cond_mel_lengths,
"do_sample": False,
"top_p": 0.8,
"top_k": None,
"temperature": 1.0,
"num_return_sequences": 1,
"length_penalty": 0.0,
"num_beams": 1,
"repetition_penalty": 10.0,
"max_generate_length": 100,
}
# baseline for non-pad
baseline = tts.gpt.inference_speech(auto_conditioning, text_tokens, **kwargs)
baseline = baseline.squeeze(0)
print("Inference padded text tokens...")
pad_text_tokens = [
F.pad(text_tokens, (8, 0), value=0), # left bos
F.pad(text_tokens, (0, 8), value=1), # right eos
F.pad(F.pad(text_tokens, (4, 0), value=0), (0, 4), value=1), # both side
F.pad(F.pad(text_tokens, (6, 0), value=0), (0, 2), value=1),
F.pad(F.pad(text_tokens, (0, 4), value=0), (0, 4), value=1),
]
output_for_padded = []
for t in pad_text_tokens:
# test for each padded text
out = tts.gpt.inference_speech(auto_conditioning, text_tokens, **kwargs)
output_for_padded.append(out.squeeze(0))
# batched inference
print("Inference padded text tokens as one batch...")
batched_text_tokens = torch.cat(pad_text_tokens, dim=0).to(tts.device)
assert len(pad_text_tokens) == batched_text_tokens.shape[0] and batched_text_tokens.ndim == 2
batch_output = tts.gpt.inference_speech(auto_conditioning, batched_text_tokens, **kwargs)
del pad_text_tokens
mismatch_idx = []
print("baseline:", baseline.shape, baseline)
print("--"*10)
print("baseline vs padded output:")
for i in range(len(output_for_padded)):
if not baseline.equal(output_for_padded[i]):
mismatch_idx.append(i)
if len(mismatch_idx) > 0:
print("mismatch:", mismatch_idx)
for i in mismatch_idx:
print(f"[{i}]: {output_for_padded[i]}")
else:
print("all matched")
del output_for_padded
print("--"*10)
print("baseline vs batched output:")
mismatch_idx = []
for i in range(batch_output.shape[0]):
if not baseline.equal(batch_output[i]):
mismatch_idx.append(i)
if len(mismatch_idx) > 0:
print("mismatch:", mismatch_idx)
for i in mismatch_idx:
print(f"[{i}]: {batch_output[i]}")
else:
print("all matched")
print("Test finished.")