Update modeling.py

modeling.py

@@ -10,7 +10,7 @@ import argparse
 from model import OptimizedAudioRestorationModel
 import librosa
 from inference_long import apply_overlap_windowing_waveform, reconstruct_waveform_from_windows
-from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -46,8 +46,8 @@ class VoiceRestore(PreTrainedModel):
 
         # Optimized restoration model
         self.optimized_model = OptimizedAudioRestorationModel(device=device, bigvgan_model=self.bigvgan_model)
-        save_path = "/content/voicerestore/checkpoints/voice-restore-20d-16h-optim.pt"
-        state_dict = torch.load(save_path, map_location=torch.device(device))
+        save_path = "./pytorch_model.safetensors"
+        state_dict = load_file(save_path, device=device)
         if 'model_state_dict' in state_dict:
             state_dict = state_dict['model_state_dict']