diff --git a/indextts/gpt/model.py b/indextts/gpt/model.py index e432c00..6f98a45 100644 --- a/indextts/gpt/model.py +++ b/indextts/gpt/model.py @@ -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] - ] - ``` - text_masks: [b, t] - padded_masks: [b, t+2] + 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, 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) + 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, + ] + # +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 diff --git a/indextts/infer.py b/indextts/infer.py index 37bda7f..3d07dfe 100644 --- a/indextts/infer.py +++ b/indextts/infer.py @@ -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) - print(f"wav shape: {wav.shape}", "min:", wav.min(), "max:", wav.max()) + 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") diff --git a/tests/padding_test.py b/tests/padding_test.py new file mode 100644 index 0000000..9fe418b --- /dev/null +++ b/tests/padding_test.py @@ -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.") \ No newline at end of file