Update model.py

model.py

@@ -43,9 +43,9 @@ class INF5Model(PreTrainedModel):
         self.vocoder = torch.compile(load_vocoder(vocoder_name="vocos", is_local=False, device=device))
 
         # Download and load model weights
-        # safetensors_path = hf_hub_download(config.name_or_path, filename="model.safetensors")
-        # print(f"Loading model weights from {safetensors_path} (safetensors)...")
-        # state_dict = load_file(safetensors_path, device=str(device))
+        safetensors_path = hf_hub_download(config.name_or_path, filename="model.safetensors")
+        print(f"Loading model weights from {safetensors_path} (safetensors)...")
+        state_dict = load_file(safetensors_path, device=str(device))
 
         # Download vocab.txt from HF Hub
         vocab_path = hf_hub_download(config.name_or_path, filename="checkpoints/vocab.txt")
@@ -60,7 +60,7 @@ class INF5Model(PreTrainedModel):
         )
 
         # # Load state dict into model
-        # self.ema_model.load_state_dict(state_dict, strict=False)
+        self.ema_model.load_state_dict(state_dict, strict=False)
 
     
     def forward(self, text: str, ref_audio_path: str, ref_text: str):