mirror of
https://github.com/Anjok07/ultimatevocalremovergui.git
synced 2024-11-14 19:07:37 +01:00
108 lines
3.1 KiB
Python
108 lines
3.1 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
|
|
|
|
from diffq import DiffQuantizer
|
|
import torch.hub
|
|
|
|
from .model import Demucs
|
|
from .tasnet import ConvTasNet
|
|
from .utils import set_state
|
|
|
|
logger = logging.getLogger(__name__)
|
|
ROOT = "https://dl.fbaipublicfiles.com/demucs/v3.0/"
|
|
|
|
PRETRAINED_MODELS = {
|
|
'demucs': 'e07c671f',
|
|
'demucs48_hq': '28a1282c',
|
|
'demucs_extra': '3646af93',
|
|
'demucs_quantized': '07afea75',
|
|
'tasnet': 'beb46fac',
|
|
'tasnet_extra': 'df3777b2',
|
|
'demucs_unittest': '09ebc15f',
|
|
}
|
|
|
|
SOURCES = ["drums", "bass", "other", "vocals"]
|
|
|
|
|
|
def get_url(name):
|
|
sig = PRETRAINED_MODELS[name]
|
|
return ROOT + name + "-" + sig[:8] + ".th"
|
|
|
|
|
|
def is_pretrained(name):
|
|
return name in PRETRAINED_MODELS
|
|
|
|
|
|
def load_pretrained(name):
|
|
if name == "demucs":
|
|
return demucs(pretrained=True)
|
|
elif name == "demucs48_hq":
|
|
return demucs(pretrained=True, hq=True, channels=48)
|
|
elif name == "demucs_extra":
|
|
return demucs(pretrained=True, extra=True)
|
|
elif name == "demucs_quantized":
|
|
return demucs(pretrained=True, quantized=True)
|
|
elif name == "demucs_unittest":
|
|
return demucs_unittest(pretrained=True)
|
|
elif name == "tasnet":
|
|
return tasnet(pretrained=True)
|
|
elif name == "tasnet_extra":
|
|
return tasnet(pretrained=True, extra=True)
|
|
else:
|
|
raise ValueError(f"Invalid pretrained name {name}")
|
|
|
|
|
|
def _load_state(name, model, quantizer=None):
|
|
url = get_url(name)
|
|
state = torch.hub.load_state_dict_from_url(url, map_location='cpu', check_hash=True)
|
|
set_state(model, quantizer, state)
|
|
if quantizer:
|
|
quantizer.detach()
|
|
|
|
|
|
def demucs_unittest(pretrained=True):
|
|
model = Demucs(channels=4, sources=SOURCES)
|
|
if pretrained:
|
|
_load_state('demucs_unittest', model)
|
|
return model
|
|
|
|
|
|
def demucs(pretrained=True, extra=False, quantized=False, hq=False, channels=64):
|
|
if not pretrained and (extra or quantized or hq):
|
|
raise ValueError("if extra or quantized is True, pretrained must be True.")
|
|
model = Demucs(sources=SOURCES, channels=channels)
|
|
if pretrained:
|
|
name = 'demucs'
|
|
if channels != 64:
|
|
name += str(channels)
|
|
quantizer = None
|
|
if sum([extra, quantized, hq]) > 1:
|
|
raise ValueError("Only one of extra, quantized, hq, can be True.")
|
|
if quantized:
|
|
quantizer = DiffQuantizer(model, group_size=8, min_size=1)
|
|
name += '_quantized'
|
|
if extra:
|
|
name += '_extra'
|
|
if hq:
|
|
name += '_hq'
|
|
_load_state(name, model, quantizer)
|
|
return model
|
|
|
|
|
|
def tasnet(pretrained=True, extra=False):
|
|
if not pretrained and extra:
|
|
raise ValueError("if extra is True, pretrained must be True.")
|
|
model = ConvTasNet(X=10, sources=SOURCES)
|
|
if pretrained:
|
|
name = 'tasnet'
|
|
if extra:
|
|
name = 'tasnet_extra'
|
|
_load_state(name, model)
|
|
return model
|