From d44d721d6920297a0d3d06bd27c3a372613a80bb Mon Sep 17 00:00:00 2001 From: IIEleven11 Date: Wed, 31 Jul 2024 18:30:47 -0700 Subject: [PATCH] Tokenizer changes --- finetune.py | 32 ++++-- start_alltalk.sh | 4 + start_diagnostics.sh | 4 + start_environment.sh | 16 +++ start_finetune.sh | 5 + system/ft_tokenizer/compare_and_merge.py | 48 +++++++++ system/ft_tokenizer/expand_xtts.py | 126 +++++++++++++++++++++++ 7 files changed, 227 insertions(+), 8 deletions(-) create mode 100755 start_alltalk.sh create mode 100755 start_diagnostics.sh create mode 100755 start_environment.sh create mode 100755 start_finetune.sh create mode 100644 system/ft_tokenizer/compare_and_merge.py create mode 100644 system/ft_tokenizer/expand_xtts.py diff --git a/finetune.py b/finetune.py index f4b11f8..073bf71 100644 --- a/finetune.py +++ b/finetune.py @@ -2376,12 +2376,28 @@ if __name__ == "__main__": ) model_to_train.change(basemodel_or_finetunedmodel_choice, model_to_train, None) - demo.queue().launch( - show_api=False, - inbrowser=True, - share=False, - debug=False, - server_port=7052, - server_name="127.0.0.1", - ) +import random +def find_available_port(start_port, end_port): + ports = list(range(start_port, end_port + 1)) + random.shuffle(ports) + return ports + +ports_to_try = find_available_port(7800, 7810) + +for port in ports_to_try: + try: + demo.queue().launch( + show_api=False, + inbrowser=True, + share=False, + debug=False, + server_port=port, + server_name="127.0.0.1", + ) + print(f"Successfully launched on port {port}") + break + except OSError: + print(f"Port {port} is not available, trying next...") +else: + print("No available ports in the specified range.") diff --git a/start_alltalk.sh b/start_alltalk.sh new file mode 100755 index 0000000..55e0756 --- /dev/null +++ b/start_alltalk.sh @@ -0,0 +1,4 @@ +#!/bin/bash +source "/home/eleven/alltalkbeta/alltalk_tts/alltalk_environment/conda/etc/profile.d/conda.sh" +conda activate "/home/eleven/alltalkbeta/alltalk_tts/alltalk_environment/env" +python script.py diff --git a/start_diagnostics.sh b/start_diagnostics.sh new file mode 100755 index 0000000..c558687 --- /dev/null +++ b/start_diagnostics.sh @@ -0,0 +1,4 @@ +#!/bin/bash +source "/home/eleven/alltalkbeta/alltalk_tts/alltalk_environment/conda/etc/profile.d/conda.sh" +conda activate "/home/eleven/alltalkbeta/alltalk_tts/alltalk_environment/env" +python diagnostics.py diff --git a/start_environment.sh b/start_environment.sh new file mode 100755 index 0000000..9791ff3 --- /dev/null +++ b/start_environment.sh @@ -0,0 +1,16 @@ +#!/bin/bash +cd "." +if [[ "/home/eleven/alltalkbeta/alltalk_tts" =~ " " ]]; then echo This script relies on Miniconda which can not be silently installed under a path with spaces. && exit; fi +# deactivate existing conda envs as needed to avoid conflicts +{ conda deactivate && conda deactivate && conda deactivate; } 2> /dev/null +# config +CONDA_ROOT_PREFIX="/home/eleven/alltalkbeta/alltalk_tts/alltalk_environment/conda" +INSTALL_ENV_DIR="/home/eleven/alltalkbeta/alltalk_tts/alltalk_environment/env" +# environment isolation +export PYTHONNOUSERSITE=1 +unset PYTHONPATH +unset PYTHONHOME +export CUDA_PATH="/home/eleven/alltalkbeta/alltalk_tts/alltalk_environment/env" +export CUDA_HOME="/home/eleven/alltalkbeta/alltalk_tts/alltalk_environment/env" +# activate env +bash --init-file <(echo "source \"/home/eleven/alltalkbeta/alltalk_tts/alltalk_environment/conda/etc/profile.d/conda.sh\" && conda activate \"/home/eleven/alltalkbeta/alltalk_tts/alltalk_environment/env\"") diff --git a/start_finetune.sh b/start_finetune.sh new file mode 100755 index 0000000..c09d800 --- /dev/null +++ b/start_finetune.sh @@ -0,0 +1,5 @@ +#!/bin/bash +export TRAINER_TELEMETRY=0 +source "/home/eleven/alltalkbeta/alltalk_tts/alltalk_environment/conda/etc/profile.d/conda.sh" +conda activate "/home/eleven/alltalkbeta/alltalk_tts/alltalk_environment/env" +python finetune.py diff --git a/system/ft_tokenizer/compare_and_merge.py b/system/ft_tokenizer/compare_and_merge.py new file mode 100644 index 0000000..dd1c453 --- /dev/null +++ b/system/ft_tokenizer/compare_and_merge.py @@ -0,0 +1,48 @@ +import json +import os + +# Define file paths +base_path = "/alltalk_tts/system/ft_tokenizer/" +original_file = os.path.join(base_path, "/alltalk_tts/models/xtts/xttsv2_2.0.3/vocab.json") # base model vocab.json +new_file = os.path.join(base_path, "/expanded_model/bpe_tokenizer-vocab.json") # new bpe tokenizer vocab.json +output_file = os.path.join(base_path, "/expanded_model/expanded_vocab.json") # output expanded_vocab.json file/path + + +with open(original_file, 'r') as f: + original_data = json.load(f) + +with open(new_file, 'r') as f: + new_data = json.load(f) + +# Get the original and new vocabularies +original_vocab = original_data['model']['vocab'] +new_vocab = new_data['model']['vocab'] + +# Find the maximum value in the original vocabulary. We do this so we make sure to append and not overwrite any new values +max_value = max(original_vocab.values()) + +# merge +for key, value in new_vocab.items(): + if key not in original_vocab: + max_value += 1 + original_vocab[key] = max_value + +# Update +original_data['model']['vocab'] = original_vocab + +# Get the original and new merges +original_merges = original_data['model']['merges'] +new_merges = new_data['model']['merges'] + +# Merge the merges +merged_merges = original_merges.copy() +for merge in new_merges: + if merge not in merged_merges: + merged_merges.append(merge) + +# Update +original_data['model']['merges'] = merged_merges +with open(output_file, 'w') as f: + json.dump(original_data, f, ensure_ascii=False, indent=2) + +print(f"Merged vocabulary and merges saved to {output_file}") diff --git a/system/ft_tokenizer/expand_xtts.py b/system/ft_tokenizer/expand_xtts.py new file mode 100644 index 0000000..2c79f0c --- /dev/null +++ b/system/ft_tokenizer/expand_xtts.py @@ -0,0 +1,126 @@ +""" +This script does not have any specific alltalk integration yet. So we are manually entering paths and it works as a standalone. + +The script does the following: + - Expands the embedding layer of the base XTTSv2 model according to the user created/trained bpe_tokenizer-vocab.json and the base model vocab.json. +Set variable paths with the base config and base model.pth. +Set the new tokenizer/vocab bpe_tokenizer-vocab.json location. +The new model will be saved at \expanded_models\expanded_model.pth + +Once this is done the new expanded model must be swapped in for the base model.pth then combined with the dvae.pth, vocab.json(bpe_tokenizer-vocab.json), base model config.json, +base model speaker_xtts.pth and base model vocoder.json. + + +I left some print debug statements in the script, they may be nice for the user to see during the process. + +""" + +import torch +import torch.nn as nn +import json +from TTS.tts.models.xtts import Xtts +from TTS.tts.layers.xtts.trainer.gpt_trainer import GPTTrainerConfig +from TTS.tts.layers.xtts.tokenizer import VoiceBpeTokenizer + +config_path = "/alltalk_tts/models/xtts/xttsv2_2.0.3/config.json" # Path to the base model config.json +pretrained_model_path = "/alltalk_tts/models/xtts/xttsv2_2.0.3/model.pth" # Path to the base model.pth +new_tokenizer_path = "/expanded_model/expanded_vocab.json" # Path to the new combined expanded_vocab.json +expanded_model_path = "/expanded_model/expanded_model.pth" # Path to where you want the new expanded_model.pth + + +# Open and load the configuration file +with open(config_path, "r") as f: + config_dict = json.load(f) + +# Create a GPTTrainerConfig object and populate it with the loaded configuration +config = GPTTrainerConfig() +config.from_dict(data=config_dict) + +# Function to get the vocabulary size from a tokenizer file +def get_vocab_size(tokenizer_path): + tokenizer = VoiceBpeTokenizer(vocab_file=tokenizer_path) + return len(tokenizer.tokenizer.get_vocab()) + +# 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 + 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 + if missing_keys: + print(f"Missing keys: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys: {unexpected_keys}") + print("Pretrained model loaded successfully.") + + # Create a new tokenizer with the new vocabulary + new_tokenizer = VoiceBpeTokenizer(vocab_file=new_tokenizer_path) + + # Get the old and new vocabulary sizes, and the embedding dimension + old_vocab_size = model.gpt.text_embedding.num_embeddings + 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 + ) + +# 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 +): + old_vocab_size = model.gpt.text_embedding.num_embeddings + embedding_dim = model.gpt.text_embedding.embedding_dim + + # Create new embedding and linear layers with the new vocabulary size + new_text_embedding = nn.Embedding(new_vocab_size, embedding_dim) + new_text_head = nn.Linear(embedding_dim, new_vocab_size) + + # 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_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_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}") + +# Expand the pretrained model with the new tokenizer +adjust_pretrained_model(pretrained_model_path, expanded_model_path, new_tokenizer_path)