|
|
from transformers import VitsModel , VitsConfig |
|
|
from torch import nn |
|
|
from torch.nn.utils.parametrizations import weight_norm |
|
|
from transformers.modeling_utils import load_state_dict |
|
|
import torch |
|
|
|
|
|
class ModVitsModel(VitsModel): |
|
|
def __init__(self, config: VitsConfig): |
|
|
config.num_speakers = len(config.emotion_names) * len(config.speaker_names) |
|
|
super().__init__(config) |
|
|
def init_weights(self): |
|
|
self.decoder.upsampler = nn.ModuleList([weight_norm(layer) for layer in self.decoder.upsampler]) |
|
|
for block in self.decoder.resblocks: |
|
|
block.convs1 = nn.ModuleList([weight_norm(layer) for layer in block.convs1]) |
|
|
block.convs2 = nn.ModuleList([weight_norm(layer) for layer in block.convs2]) |
|
|
return super().init_weights() |
|
|
@classmethod |
|
|
def _load_pretrained_model(cls, model, state_dict, checkpoint_files, pretrained_model_name_or_path, ignore_mismatched_sizes = False, sharded_metadata = None, device_map = None, disk_offload_folder = None, dtype = None, hf_quantizer = None, keep_in_fp32_regex = None, device_mesh = None, key_mapping = None, weights_only = True): |
|
|
state_dict = load_state_dict(checkpoint_files[0],weights_only=weights_only) |
|
|
emotions = state_dict['embed_emotion.weight'][:len(model.config.emotion_names)+len(model.config.undefined_emotion_index)][[i for i in range(len(model.config.emotion_names)+len(model.config.undefined_emotion_index)) if i not in model.config.undefined_emotion_index]] |
|
|
speakers = state_dict['embed_speaker.weight'][:len(model.config.speaker_names)] |
|
|
state_dict['embed_speaker.weight'] = torch.stack([s + e for s in speakers for e in emotions]).reshape(-1, model.config.speaker_embedding_size) |
|
|
del state_dict['embed_emotion.weight'] |
|
|
return super()._load_pretrained_model(model, state_dict, checkpoint_files, pretrained_model_name_or_path, ignore_mismatched_sizes, sharded_metadata, device_map, disk_offload_folder, dtype, hf_quantizer, keep_in_fp32_regex, device_mesh, key_mapping, weights_only) |
|
|
|
|
|
def forward(self, input_ids = None, attention_mask = None, speaker_id = 'Marathi-Male', emotion_id = 'ALEXA',output_attentions = None, output_hidden_states = None, return_dict = None, labels = None): |
|
|
speaker_id = self.config.speaker_names.index(speaker_id) * len(self.config.emotion_names) + self.config.emotion_names.index(emotion_id) |
|
|
return super().forward(input_ids, attention_mask, speaker_id, output_attentions, output_hidden_states, return_dict, labels) |
|
|
|
