From aafae8da82803857f3f0b6e289d22d552eeae2d1 Mon Sep 17 00:00:00 2001 From: IIEleven11 Date: Thu, 19 Sep 2024 21:11:31 -0700 Subject: [PATCH] Fixed slurred speech by freezing all layers but the two embedding layers --- system/ft_tokenizer/expand_xtts.py | 51 ++++++++++++++---------------- 1 file changed, 23 insertions(+), 28 deletions(-) diff --git a/system/ft_tokenizer/expand_xtts.py b/system/ft_tokenizer/expand_xtts.py index 2c79f0c..832fa7b 100644 --- a/system/ft_tokenizer/expand_xtts.py +++ b/system/ft_tokenizer/expand_xtts.py @@ -43,20 +43,13 @@ def get_vocab_size(tokenizer_path): # Function to adjust the pretrained model with a new tokenizer def adjust_pretrained_model( - pretrained_model_path, adjusted_model_path, new_tokenizer_path -): - # Load the pretrained model state dictionary + pretrained_model_path, adjusted_model_path, new_tokenizer_path): state_dict = torch.load(pretrained_model_path) pretrained_state_dict = state_dict["model"] - - # Create a new Xtts model instance with the loaded configuration model = Xtts(config) # Load the pretrained state dictionary into the new model - missing_keys, unexpected_keys = model.load_state_dict( - pretrained_state_dict, strict=False - ) - # Print any missing or unexpected keys for debugging + missing_keys, unexpected_keys = model.load_state_dict(pretrained_state_dict, strict=False) if missing_keys: print(f"Missing keys: {missing_keys}") if unexpected_keys: @@ -71,20 +64,18 @@ def adjust_pretrained_model( new_vocab_size = len(new_tokenizer.tokenizer.get_vocab()) embedding_dim = model.gpt.text_embedding.embedding_dim - # Print vocabulary sizes and embedding dimension for debugging print(f"Old vocab size: {old_vocab_size}") print(f"New vocab size: {new_vocab_size}") print(f"Embedding dimension: {embedding_dim}") # Adjust the embedding layer with the new vocabulary size - adjust_embedding_layer( - model, pretrained_state_dict, new_vocab_size, adjusted_model_path - ) + adjust_embedding_layer(model, new_vocab_size, adjusted_model_path) -# Function to adjust the embedding layer for the new vocabulary size -def adjust_embedding_layer( - model, pretrained_state_dict, new_vocab_size, adjusted_model_path -): + # Freeze all parameters except the position embeddings + freeze_except_position_embeddings(model) + + # Function to adjust the embedding layer for the new vocabulary size +def adjust_embedding_layer(model, new_vocab_size, adjusted_model_path): old_vocab_size = model.gpt.text_embedding.num_embeddings embedding_dim = model.gpt.text_embedding.embedding_dim @@ -94,33 +85,37 @@ def adjust_embedding_layer( # Copy weights from the old embedding layer to the new one if new_vocab_size > old_vocab_size: - # If the new vocabulary is larger, copy existing weights and initialize new ones - new_text_embedding.weight.data[:old_vocab_size] = ( - model.gpt.text_embedding.weight.data - ) + new_text_embedding.weight.data[:old_vocab_size] = model.gpt.text_embedding.weight.data new_text_head.weight.data[:old_vocab_size] = model.gpt.text_head.weight.data new_text_head.bias.data[:old_vocab_size] = model.gpt.text_head.bias.data - # Initialize new weights with normal distribution new_text_embedding.weight.data[old_vocab_size:].normal_(mean=0.0, std=0.02) new_text_head.weight.data[old_vocab_size:].normal_(mean=0.0, std=0.02) new_text_head.bias.data[old_vocab_size:].normal_(mean=0.0, std=0.02) else: - # If the new vocabulary is smaller, truncate the existing weights - new_text_embedding.weight.data = model.gpt.text_embedding.weight.data[ - :new_vocab_size - ] + new_text_embedding.weight.data = model.gpt.text_embedding.weight.data[:new_vocab_size] new_text_head.weight.data = model.gpt.text_head.weight.data[:new_vocab_size] new_text_head.bias.data = model.gpt.text_head.bias.data[:new_vocab_size] - # Replace the old embedding layer with the new one model.gpt.text_embedding = new_text_embedding model.gpt.text_head = new_text_head - # Save the adjusted model checkpoint = {"model": model.state_dict()} torch.save(checkpoint, adjusted_model_path) print(f"Adjusted model saved to {adjusted_model_path}") +# Function to freeze all parameters except the position embeddings +def freeze_except_position_embeddings(model): + for param in model.parameters(): + param.requires_grad = False + + for name, param in model.named_parameters(): + if 'pos_embedding' in name: + param.requires_grad = True + + # Verify which parameters are frozen and which are not, comment this out if you dont want to debug. You should see only two true values + for name, param in model.named_parameters(): + print(f"{name}: requires_grad={param.requires_grad}") + # Expand the pretrained model with the new tokenizer adjust_pretrained_model(pretrained_model_path, expanded_model_path, new_tokenizer_path)