Tran Dang Khoa
Add application file
cafad09
Raw
History Blame Contribute Delete
1.92 kB
"""Model construction helpers."""
from __future__ import annotations
from .config import load_config
def get_model_from_config(model_type, config_path, model_kwargs_override=None):
"""Instantiate a separation model from a model configuration file."""
model_kwargs_override = model_kwargs_override or {}
config = load_config(config_path)
if model_type == "mdx23c":
from .modules.mdx23c_tfc_tdf_v3 import TFC_TDF_net
return TFC_TDF_net(config), config
if model_type == "htdemucs":
from .modules.demucs4ht import get_model
return get_model(config), config
if model_type == "mel_band_roformer":
from .modules.bs_roformer import MelBandRoformer
model_kwargs = dict(config.model)
model_kwargs.update(model_kwargs_override)
return MelBandRoformer(**model_kwargs), config
if model_type == "bs_roformer":
from .modules.bs_roformer import BSRoformer
return BSRoformer(**dict(config.model)), config
if model_type == "bs_roformer_hyperace":
from .modules.bs_roformer import BSRoformerHyperACE
return BSRoformerHyperACE(**dict(config.model)), config
if model_type == "bandit":
from .modules.bandit.core.model import MultiMaskMultiSourceBandSplitRNNSimple
return MultiMaskMultiSourceBandSplitRNNSimple(**config.model), config
if model_type == "bandit_v2":
from .modules.bandit_v2.bandit import Bandit
return Bandit(**config.kwargs), config
if model_type == "scnet":
from .modules.scnet import SCNet
return SCNet(**config.model), config
if model_type == "apollo":
from .modules.look2hear.apollo import Apollo
return Apollo(**config.model), config
if model_type == "vr":
raise ValueError("VR network modules do not use YAML config loading")
raise ValueError(f"Model type {model_type} not supported")