122 lines
5.7 KiB
Python
122 lines
5.7 KiB
Python
"""
|
|
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):
|
|
state_dict = torch.load(pretrained_model_path)
|
|
pretrained_state_dict = state_dict["model"]
|
|
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)
|
|
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(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, 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
|
|
|
|
# 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:
|
|
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
|
|
|
|
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:
|
|
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]
|
|
|
|
model.gpt.text_embedding = new_text_embedding
|
|
model.gpt.text_head = new_text_head
|
|
|
|
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)
|