Update modeling_vits.py

modeling_vits.py

@@ -14,16 +14,16 @@ class ModVitsModel(VitsModel):
           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)
+
+    @staticmethod
+    def _load_pretrained_model(model, state_dict, checkpoint_files, load_config):
+        state_dict = load_state_dict(checkpoint_files[0])
         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)
+        return super()._load_pretrained_model(model, state_dict, checkpoint_files, load_config)
 
     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)
-