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)