48 lines
1.4 KiB
Python
48 lines
1.4 KiB
Python
from pathlib import Path
|
|
import typing as tp
|
|
|
|
from .repo import AnyModelRepo, BagOnlyRepo,RemoteRepo, ModelOnlyRepo
|
|
from .states import _check_diffq
|
|
|
|
ROOT_URL = "https://dl.fbaipublicfiles.com/demucs/"
|
|
REMOTE_ROOT = Path(__file__).parent / 'remote'
|
|
|
|
SOURCES = ["drums", "bass", "other", "vocals"]
|
|
DEFAULT_MODEL = 'htdemucs'
|
|
|
|
|
|
def _parse_remote_files(remote_file_list) -> tp.Dict[str, str]:
|
|
root: str = ''
|
|
models: tp.Dict[str, str] = {}
|
|
for line in remote_file_list.read_text().split('\n'):
|
|
line = line.strip()
|
|
if line.startswith('#'):
|
|
continue
|
|
elif len(line) == 0:
|
|
continue
|
|
elif line.startswith('root:'):
|
|
root = line.split(':', 1)[1].strip()
|
|
else:
|
|
sig = line.split('-', 1)[0]
|
|
assert sig not in models
|
|
models[sig] = ROOT_URL + root + line
|
|
return models
|
|
|
|
def get_model(name: str,
|
|
repo: tp.Optional[Path] = None):
|
|
|
|
model_repo: ModelOnlyRepo
|
|
if repo is None:
|
|
models = _parse_remote_files(REMOTE_ROOT / 'files.txt')
|
|
model_repo = RemoteRepo(models)
|
|
bag_repo = BagOnlyRepo(REMOTE_ROOT, model_repo)
|
|
any_repo = AnyModelRepo(model_repo, bag_repo)
|
|
try:
|
|
model = any_repo.get_model(name)
|
|
except ImportError as exc:
|
|
if 'diffq' in exc.args[0]:
|
|
_check_diffq()
|
|
raise
|
|
|
|
model.eval()
|
|
return model |