# ported from: https://github.com/neonbjb/tortoise-tts # ported from: https://github.com/coqui-ai/TTS/blob/dev/TTS/tts/layers/xtts/gpt.py import functools import math import random import torch import torch.nn as nn import torch.nn.functional as F from transformers import GPT2Config, GPT2Model, GPT2PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions class GPT2InferenceModel(GPT2PreTrainedModel): """Override GPT2LMHeadModel to allow for prefix conditioning.""" def __init__(self, config, gpt, pos_emb, embeddings, norm, linear, kv_cache): super().__init__(config) self.transformer = gpt self.pos_embedding = pos_emb self.embeddings = embeddings self.final_norm = norm self.lm_head = nn.Sequential(norm, linear) self.kv_cache = kv_cache def store_prefix_emb(self, prefix_emb): self.cached_prefix_emb = prefix_emb def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs): token_type_ids = kwargs.get("token_type_ids", None) # usually None if not self.kv_cache: past_key_values = None # only last token for inputs_ids if past is defined in kwargs if past_key_values is not None: input_ids = input_ids[:, -1].unsqueeze(-1) if token_type_ids is not None: token_type_ids = token_type_ids[:, -1].unsqueeze(-1) attention_mask = kwargs.get("attention_mask", None) position_ids = kwargs.get("position_ids", None) 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) if past_key_values is not None: position_ids = position_ids[:, -1].unsqueeze(-1) else: position_ids = None return { "input_ids": input_ids, "past_key_values": past_key_values, "use_cache": kwargs.get("use_cache"), "position_ids": position_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids, } def forward( self, input_ids=None, past_key_values=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, encoder_hidden_states=None, encoder_attention_mask=None, labels=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None, ): assert self.cached_prefix_emb is not None assert inputs_embeds is None # Not supported by this inference model. assert labels is None # Training not supported by this inference model. return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) # Create embedding prefix_len = self.cached_prefix_emb.shape[1] if input_ids.shape[1] != 1: gen_inputs = input_ids[:, prefix_len:] gen_emb = self.embeddings(gen_inputs) gen_emb = gen_emb + self.pos_embedding(gen_emb) if self.cached_prefix_emb.shape[0] != gen_emb.shape[0]: prefix_emb = self.cached_prefix_emb.repeat_interleave( gen_emb.shape[0] // self.cached_prefix_emb.shape[0], 0 ) else: prefix_emb = self.cached_prefix_emb.to(gen_emb.dtype) emb = torch.cat([prefix_emb, gen_emb], dim=1) else: emb = self.embeddings(input_ids) emb = emb + self.pos_embedding.get_fixed_embedding( attention_mask.shape[1] - (prefix_len + 1), attention_mask.device ) transformer_outputs = self.transformer( inputs_embeds=emb, past_key_values=past_key_values, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) hidden_states = transformer_outputs[0] lm_logits = self.lm_head(hidden_states) if not return_dict: return (lm_logits,) + transformer_outputs[1:] return CausalLMOutputWithCrossAttentions( loss=None, logits=lm_logits, past_key_values=transformer_outputs.past_key_values, hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions, cross_attentions=transformer_outputs.cross_attentions, ) @staticmethod def _reorder_cache(past, beam_idx): """ This function is used to re-order the :obj:`past_key_values` cache if :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step. """ return tuple( tuple( past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past ) for layer_past in past ) def null_position_embeddings(range, dim): return torch.zeros((range.shape[0], range.shape[1], dim), device=range.device) class LearnedPositionEmbeddings(nn.Module): def __init__(self, seq_len, model_dim, init=0.02): super().__init__() self.emb = torch.nn.Embedding(seq_len, model_dim) # Initializing this way is standard for GPT-2 self.emb.weight.data.normal_(mean=0.0, std=init) def forward(self, x): sl = x.shape[1] return self.emb(torch.arange(0, sl, device=x.device)) def get_fixed_embedding(self, ind, dev): return self.emb(torch.tensor([ind], device=dev)).unsqueeze(0) def build_hf_gpt_transformer( layers, model_dim, heads, max_mel_seq_len, max_text_seq_len, max_prompt_len, checkpointing, ): """ GPT-2 implemented by the HuggingFace library. """ gpt_config = GPT2Config( vocab_size=256, # Unused. n_positions=max_mel_seq_len + max_text_seq_len + max_prompt_len, n_ctx=max_mel_seq_len + max_text_seq_len + max_prompt_len, n_embd=model_dim, n_layer=layers, n_head=heads, gradient_checkpointing=checkpointing, use_cache=not checkpointing, ) gpt = GPT2Model(gpt_config) # Override the built in positional embeddings del gpt.wpe gpt.wpe = functools.partial(null_position_embeddings, dim=model_dim) # Built-in token embeddings are unused. del gpt.wte mel_pos_emb = ( LearnedPositionEmbeddings(max_mel_seq_len, model_dim) if max_mel_seq_len != -1 else functools.partial(null_position_embeddings, dim=model_dim) ) text_pos_emb = ( LearnedPositionEmbeddings(max_text_seq_len, model_dim) if max_mel_seq_len != -1 else functools.partial(null_position_embeddings, dim=model_dim) ) return gpt, mel_pos_emb, text_pos_emb, None, None class GPT(nn.Module): def __init__( self, start_text_token=261, stop_text_token=0, layers=8, model_dim=512, heads=8, max_text_tokens=120, max_mel_tokens=250, max_prompt_tokens=70, max_conditioning_inputs=1, code_stride_len=1024, number_text_tokens=256, num_audio_tokens=8194, start_audio_token=8192, stop_audio_token=8193, checkpointing=False, label_smoothing=0.0, ): """ Args: """ super().__init__() self.label_smoothing = label_smoothing self.number_text_tokens = number_text_tokens self.start_text_token = start_text_token self.stop_text_token = stop_text_token self.num_audio_tokens = num_audio_tokens self.start_audio_token = start_audio_token self.stop_audio_token = stop_audio_token self.start_prompt_token = start_audio_token self.stop_prompt_token = stop_audio_token self.layers = layers self.heads = heads self.model_dim = model_dim self.max_conditioning_inputs = max_conditioning_inputs self.max_gen_mel_tokens = max_mel_tokens - self.max_conditioning_inputs - 2 self.max_mel_tokens = ( -1 if max_mel_tokens == -1 else max_mel_tokens + 2 + self.max_conditioning_inputs ) self.max_text_tokens = -1 if max_text_tokens == -1 else max_text_tokens + 2 self.max_prompt_tokens = max_prompt_tokens self.code_stride_len = code_stride_len self.conditioning_dropout = nn.Dropout1d(0.1) self.text_embedding = nn.Embedding(self.number_text_tokens, model_dim) self.mel_embedding = nn.Embedding(self.num_audio_tokens, model_dim) ( self.gpt, self.mel_pos_embedding, self.text_pos_embedding, self.mel_layer_pos_embedding, self.text_layer_pos_embedding, ) = build_hf_gpt_transformer( layers, model_dim, heads, self.max_mel_tokens, self.max_text_tokens, self.max_prompt_tokens, checkpointing, ) self.final_norm = nn.LayerNorm(model_dim) self.text_head = nn.Linear(model_dim, self.number_text_tokens) self.mel_head = nn.Linear(model_dim, self.num_audio_tokens) # reference_embedding self.reference_embedding = nn.Sequential( nn.Linear(512, 256), nn.Tanh(), nn.Linear(256, self.model_dim), ) def init_gpt_for_inference(self, kv_cache=True, use_deepspeed=False): seq_length = ( self.max_prompt_tokens + self.max_mel_tokens + self.max_text_tokens + 1 ) gpt_config = GPT2Config( vocab_size=self.max_mel_tokens, n_positions=seq_length, n_ctx=seq_length, n_embd=self.model_dim, n_layer=self.layers, n_head=self.heads, gradient_checkpointing=False, use_cache=True, ) self.gpt_inference = GPT2InferenceModel( gpt_config, self.gpt, self.mel_pos_embedding, self.mel_embedding, self.final_norm, self.mel_head, kv_cache=kv_cache, ) self.gpt.wte = self.mel_embedding def inference(self, cond_latents, text_inputs, **hf_generate_kwargs): self.compute_embeddings(cond_latents, text_inputs) return self.generate(cond_latents, text_inputs, **hf_generate_kwargs) def compute_embeddings( self, cond_latents, text_inputs, ): text_inputs = F.pad(text_inputs, (0, 1), value=self.stop_text_token) text_inputs = F.pad(text_inputs, (1, 0), value=self.start_text_token) emb = self.text_embedding(text_inputs) + self.text_pos_embedding(text_inputs) emb = torch.cat([cond_latents, emb], dim=1) self.gpt_inference.store_prefix_emb(emb) gpt_inputs = torch.full( ( emb.shape[0], emb.shape[1] + 1, # +1 for the start_audio_token ), fill_value=1, dtype=torch.long, device=text_inputs.device, ) gpt_inputs[:, -1] = self.start_audio_token return gpt_inputs def generate( self, cond_latents, text_inputs, **hf_generate_kwargs, ): gpt_inputs = self.compute_embeddings(cond_latents, text_inputs) gen = self.gpt_inference.generate( gpt_inputs, bos_token_id=self.start_audio_token, pad_token_id=self.stop_audio_token, eos_token_id=self.stop_audio_token, max_length=self.max_gen_mel_tokens + gpt_inputs.shape[-1], **hf_generate_kwargs, ) if "return_dict_in_generate" in hf_generate_kwargs: return gen.sequences[:, gpt_inputs.shape[1] :], gen return gen[:, gpt_inputs.shape[1] :]