22 lines
727 B
Python
22 lines
727 B
Python
# Copyright (c) 2025 Salvador E. Tropea
|
|
# Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
|
|
# License: GPLv3
|
|
# Project: ComfyUI-AudioSeparation
|
|
#
|
|
# Helper to get a model from the correct class
|
|
import logging
|
|
from .MDX_Net import MDX_Net
|
|
from ..utils.misc import NODES_NAME
|
|
|
|
logger = logging.getLogger(f"{NODES_NAME}.get_model")
|
|
|
|
|
|
# Currently we have just one type of networks, but this is a clean way to support more, or even test replacements
|
|
def get_model(d):
|
|
model_t = d['model_t'].lower()
|
|
if model_t != "mdx":
|
|
msg = f"Unknown model type `{model_t}`"
|
|
logger.error(msg)
|
|
raise ValueError(msg)
|
|
return MDX_Net(dim_f=d['mdx_dim_f_set'], ch=d['channels'], num_stages=d['stages'])
|