wehub-resource-sync
36 lines
821 B
Python
36 lines
821 B
Python
import pathlib
|
|
|
|
import torch
|
|
import yaml
|
|
|
|
from .nets import CascadedNet
|
|
|
|
|
|
class DotDict(dict):
|
|
def __getattr__(*args):
|
|
val = dict.get(*args)
|
|
return DotDict(val) if type(val) is dict else val
|
|
|
|
__setattr__ = dict.__setitem__
|
|
__delattr__ = dict.__delitem__
|
|
|
|
|
|
def load_sep_model(model_path, device='cpu'):
|
|
model_path = pathlib.Path(model_path)
|
|
config_file = model_path.with_name('config.yaml')
|
|
with open(config_file, "r") as config:
|
|
args = yaml.safe_load(config)
|
|
args = DotDict(args)
|
|
model = CascadedNet(
|
|
args.n_fft,
|
|
args.hop_length,
|
|
args.n_out,
|
|
args.n_out_lstm,
|
|
True,
|
|
is_mono=args.is_mono
|
|
)
|
|
model.to(device)
|
|
model.load_state_dict(torch.load(model_path, map_location='cpu'))
|
|
model.eval()
|
|
return model
|