wehub-resource-sync
82 lines
2.4 KiB
Python
82 lines
2.4 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
# author: adefossez
|
|
|
|
import logging
|
|
|
|
import torch.hub
|
|
|
|
from .demucs import Demucs
|
|
from .utils import deserialize_model
|
|
|
|
logger = logging.getLogger(__name__)
|
|
ROOT = "https://dl.fbaipublicfiles.com/adiyoss/denoiser/"
|
|
DNS_48_URL = ROOT + "dns48-11decc9d8e3f0998.th"
|
|
DNS_64_URL = ROOT + "dns64-a7761ff99a7d5bb6.th"
|
|
MASTER_64_URL = ROOT + "master64-8a5dfb4bb92753dd.th"
|
|
|
|
|
|
def _demucs(pretrained, url, **kwargs):
|
|
model = Demucs(**kwargs)
|
|
if pretrained:
|
|
state_dict = torch.hub.load_state_dict_from_url(url, map_location='cpu')
|
|
model.load_state_dict(state_dict)
|
|
return model
|
|
|
|
|
|
def dns48(pretrained=True):
|
|
return _demucs(pretrained, DNS_48_URL, hidden=48)
|
|
|
|
|
|
def dns64(pretrained=True):
|
|
return _demucs(pretrained, DNS_64_URL, hidden=64)
|
|
|
|
|
|
def master64(pretrained=True):
|
|
return _demucs(pretrained, MASTER_64_URL, hidden=64)
|
|
|
|
|
|
def add_model_flags(parser):
|
|
group = parser.add_mutually_exclusive_group(required=False)
|
|
group.add_argument(
|
|
"-m", "--model_path", help="Path to local trained model."
|
|
)
|
|
group.add_argument(
|
|
"--dns48", action="store_true",
|
|
help="Use pre-trained real time H=48 model trained on DNS."
|
|
)
|
|
group.add_argument(
|
|
"--dns64", action="store_true",
|
|
help="Use pre-trained real time H=64 model trained on DNS."
|
|
)
|
|
group.add_argument(
|
|
"--master64", action="store_true",
|
|
help="Use pre-trained real time H=64 model trained on DNS and Valentini."
|
|
)
|
|
|
|
|
|
def get_model(args):
|
|
"""
|
|
Load local model package or torchhub pre-trained model.
|
|
"""
|
|
if args.model_path:
|
|
logger.info("Loading model from %s", args.model_path)
|
|
pkg = torch.load(args.model_path)
|
|
model = deserialize_model(pkg)
|
|
elif args.dns64:
|
|
logger.info("Loading pre-trained real time H=64 model trained on DNS.")
|
|
model = dns64()
|
|
elif args.master64:
|
|
logger.info(
|
|
"Loading pre-trained real time H=64 model trained on DNS and Valentini."
|
|
)
|
|
model = master64()
|
|
else:
|
|
logger.info("Loading pre-trained real time H=48 model trained on DNS.")
|
|
model = dns48()
|
|
logger.debug(model)
|
|
return model
|