Update modeling_vits.py

modeling_vits.py

@@ -1,7 +1,7 @@
 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
+from safetensors.torch import load_file
 import torch
 
 class ModVitsModel(VitsModel):
@@ -17,7 +17,7 @@ class ModVitsModel(VitsModel):
 
     @staticmethod
     def _load_pretrained_model(model, state_dict, checkpoint_files, load_config):
-        state_dict = load_state_dict(checkpoint_files[0])
+        state_dict = load_file(checkpoint_files[0])
         speakers = state_dict['embed_speaker.weight'][:len(model.config.speaker_names)]
         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]]
         state_dict['embed_speaker.weight'] = torch.stack([s + e for s in speakers for e in emotions]).reshape(-1, model.config.speaker_embedding_size)
@@ -26,8 +26,10 @@ class ModVitsModel(VitsModel):
     
     @torch.inference_mode()
     def forward(self, input_ids = None, attention_mask = None, speaker_id = None, output_attentions = None, output_hidden_states = None, return_dict = None, labels = None, **kwargs):
+        speaker_id = speaker_id * len(self.config.emotion_names) + kwargs['style_id']
         audio = super().forward(input_ids, attention_mask, speaker_id, output_attentions, output_hidden_states, return_dict, labels, **kwargs)
         B, T = audio.waveform.shape
         mask = torch.arange(T, device=audio.waveform.device).expand(B, T) < audio.sequence_lengths.unsqueeze(1)
         audio.waveform.masked_fill_(~mask, 0)
         return audio
+