144 lines
5.5 KiB
Python
144 lines
5.5 KiB
Python
# Copyright (c) 2025 Salvador E. Tropea
|
|
# Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
|
|
# License: GPLv3
|
|
# Project: ComfyUI-AudioSeparation
|
|
import os
|
|
import torch
|
|
from typing import Dict
|
|
# ComfyUI imports
|
|
import folder_paths # ComfyUI's way to access model paths
|
|
# Local imports
|
|
from .source.utils.logger import main_logger
|
|
from .source.utils.load_audio import audio_get_channels, force_stereo, force_sample_rate
|
|
from .source.utils.torch import get_torch_device_options
|
|
from .source.utils.comfy_node_action import send_node_action
|
|
from .source.inference.demixer import get_demixer
|
|
from .source.db.models_db import ModelsDB
|
|
|
|
|
|
DEF_MODEL = 'Kim_Vocal_2.safetensors'
|
|
DEF_ENTRY = 'Default'
|
|
MODELS_DIR = os.path.join(folder_paths.models_dir, "audio", "MDX")
|
|
models_db = ModelsDB(MODELS_DIR)
|
|
logger = main_logger
|
|
|
|
|
|
class AudioSeparateVocals:
|
|
PRIMARY_STEM = 'Vocals'
|
|
MODEL_T = 'MDX'
|
|
FILE_T = 'safetensors'
|
|
DEFAULT_MODEL = "Kim_Vocal_2.safetensors"
|
|
|
|
@classmethod
|
|
def _get_available_audio_models(cls):
|
|
global models_db
|
|
# Refresh the database
|
|
models_db.refresh()
|
|
# Filter the models this node can handle
|
|
cls.models_filtered = models_db.get_filtered(primary_stem=cls.PRIMARY_STEM, model_t=cls.MODEL_T, file_t=cls.FILE_T,
|
|
default=cls.DEFAULT_MODEL, repeat_dl=True)
|
|
# We add any model downloaded and memorized by the GUI
|
|
return cls.models_filtered.get_display_names()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
device_options, default_device = get_torch_device_options()
|
|
return {
|
|
"required": {
|
|
"input_sound": ("AUDIO",),
|
|
"model": (cls._get_available_audio_models(),), # Dropdown for model selection
|
|
"segments": ("INT", {
|
|
"default": 1, # Default value
|
|
"min": 1, # Minimum allowed value
|
|
"max": 64, # Maximum allowed value (set a reasonable practical max)
|
|
"step": 1, # Step for slider/spinbox
|
|
"display": "slider" # How to display: "number" or "slider"
|
|
}),
|
|
"target_device": (device_options, {
|
|
"default": default_device,
|
|
"tooltip": "The device (CPU or CUDA) to which the projection layer will be assigned for computation."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("AUDIO", "AUDIO",)
|
|
RETURN_NAMES = (PRIMARY_STEM, "Complement",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "audio/separation"
|
|
DESCRIPTION = "Separates vocals using MDX-Net networks"
|
|
UNIQUE_NAME = "AudioSeparateVocals"
|
|
DISPLAY_NAME = "Vocals using MDX"
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.demixer = None
|
|
|
|
def execute(self, input_sound: Dict, model: str, segments: int, target_device: str):
|
|
# Get information for the selected model
|
|
main_logger.info(f"Selected model: {model}")
|
|
model_data = self.models_filtered.get_by_display_name(model)
|
|
if model_data is None:
|
|
raise ValueError("Unknown model selected, please refresh pressing `R` and select another")
|
|
model_path = model_data.get('model_path')
|
|
|
|
# Create or recycle a demixer
|
|
device = torch.device(target_device)
|
|
if self.demixer is None or self.demixer.d['hash'] != model_data['hash']:
|
|
# New demixer
|
|
logger.debug("Creating a new demixer object")
|
|
# This will load the model, optionally downloading it
|
|
self.demixer = get_demixer(model_data, device, MODELS_DIR)
|
|
|
|
# Handle a change in the icon of the model name
|
|
if model_path is None:
|
|
# Was downloaded
|
|
send_node_action("change_widget", "model", model_data['indicator'] + model_data['filtered_name'])
|
|
|
|
# Match channels and S/R
|
|
waveform = input_sound['waveform']
|
|
sample_rate = input_sound['sample_rate']
|
|
if audio_get_channels(waveform) == 1 and self.demixer.ch == 2:
|
|
waveform = force_stereo(waveform)
|
|
if sample_rate != self.demixer.sr:
|
|
waveform = force_sample_rate(waveform, sample_rate, self.demixer.sr)
|
|
|
|
# Demix
|
|
wavs = self.demixer(waveform, segments)
|
|
|
|
return (wavs[0], wavs[1],)
|
|
|
|
|
|
class AudioSeparateInstrumental(AudioSeparateVocals):
|
|
PRIMARY_STEM = 'Instrumental'
|
|
DEFAULT_MODEL = "Kim_Inst.safetensors"
|
|
DESCRIPTION = "Separates instruments using MDX-Net networks"
|
|
UNIQUE_NAME = "AudioSeparateInstrumental"
|
|
DISPLAY_NAME = "Instrumental using MDX"
|
|
RETURN_NAMES = (PRIMARY_STEM, "Complement",)
|
|
|
|
|
|
class AudioSeparateBass(AudioSeparateVocals):
|
|
PRIMARY_STEM = 'Bass'
|
|
DEFAULT_MODEL = "kuielab_b_bass.safetensors"
|
|
DESCRIPTION = "Separates bass using MDX-Net networks"
|
|
UNIQUE_NAME = "AudioSeparateBass"
|
|
DISPLAY_NAME = "Bass using MDX"
|
|
RETURN_NAMES = (PRIMARY_STEM, "Complement",)
|
|
|
|
|
|
class AudioSeparateDrums(AudioSeparateVocals):
|
|
PRIMARY_STEM = 'Drums'
|
|
DEFAULT_MODEL = "kuielab_b_drums.safetensors"
|
|
DESCRIPTION = "Separates drums using MDX-Net networks"
|
|
UNIQUE_NAME = "AudioSeparateDrums"
|
|
DISPLAY_NAME = "Drums using MDX"
|
|
RETURN_NAMES = (PRIMARY_STEM, "Complement",)
|
|
|
|
|
|
class AudioSeparateVarious(AudioSeparateVocals):
|
|
PRIMARY_STEM = ["Other", "Reverb"]
|
|
DEFAULT_MODEL = "Reverb_HQ_By_FoxJoy.safetensors"
|
|
DESCRIPTION = "Misc. separators using MDX-Net networks"
|
|
UNIQUE_NAME = "AudioSeparateVarious"
|
|
DISPLAY_NAME = "Various using MDX"
|
|
RETURN_NAMES = ("Main", "Complement",)
|