nvidia
/

RE-USE

universal speech enhancement

multiple input sampling rates

language-agnostic

Model card Files Files and versions

szuweifu commited on 22 days ago

Commit

2cad78e

·

verified ·

1 Parent(s): 4ffb980

Upload 2 files

Files changed (2) hide show

inference.py +1 -3
inference_chunk.py +1 -3

inference.py CHANGED Viewed

@@ -44,9 +44,7 @@ def inference(args, device):
     compress_factor = cfg['model_cfg']['compress_factor']
     sampling_rate = cfg['stft_cfg']['sampling_rate']
-    SE_model = SEMamba(cfg).to(device)
-    state_dict = torch.load(args.checkpoint_file, map_location=device)
-    SE_model.load_state_dict(state_dict['generator'])
     SE_model.eval()
     os.makedirs(args.output_folder, exist_ok=True)

     compress_factor = cfg['model_cfg']['compress_factor']
     sampling_rate = cfg['stft_cfg']['sampling_rate']
+    SE_model = SEMamba.from_pretrained("nvidia/RE-USE", cfg=cfg).to(device)
     SE_model.eval()
     os.makedirs(args.output_folder, exist_ok=True)

inference_chunk.py CHANGED Viewed

@@ -46,9 +46,7 @@ def inference(args, device):
     compress_factor = cfg['model_cfg']['compress_factor']
     sampling_rate = cfg['stft_cfg']['sampling_rate']
-    SE_model = SEMamba(cfg).to(device)
-    state_dict = torch.load(args.checkpoint_file, map_location=device)
-    SE_model.load_state_dict(state_dict['generator'])
     SE_model.eval()
     os.makedirs(args.output_folder, exist_ok=True)

     compress_factor = cfg['model_cfg']['compress_factor']
     sampling_rate = cfg['stft_cfg']['sampling_rate']
+    SE_model = SEMamba.from_pretrained("nvidia/RE-USE", cfg=cfg).to(device)
     SE_model.eval()
     os.makedirs(args.output_folder, exist_ok=True)