nodes/set-soft-AudioSeparation
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[Added] Demucs support and node

Browse Source
This commit is contained in:
Salvador E. Tropea
2025-07-06 15:09:36 -03:00
parent 0ae82e60b4
commit 2ca35cffd9
6 changed files with 513 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
source/db/models_db.py
+54 -9
View File
@@ -22,7 +22,8 @@ known_models_mtime = None
ICON_REMOTE = "⬇️ " # "\u2B07" # ⬇️
ICON_DOWNLOADED = "\U0001F4BE " # 💾 Floppy Disk
KNOWN_SOURCES = {'Politrees/MDXNet': 'https://huggingface.co/Politrees/UVR_resources/resolve/main/models/MDXNet',
'Main/MDX': 'https://huggingface.co/set-soft/audio_separation/resolve/main/MDX'}
'Main/MDX': 'https://huggingface.co/set-soft/audio_separation/resolve/main/MDX',
'Main/Demucs': 'https://huggingface.co/set-soft/audio_separation/resolve/main/Demucs', }
def get_db_filename(provided=None):
@@ -68,9 +69,10 @@ def load_known_models(json_path=None):
logger.error("Error: The models database was not found at the expected location.")
logger.error("Please check the directory structure.")
return None
except json.JSONDecodeError:
logger.error(f"Error: The file at '{json_path}' is not a valid JSON file.")
return None
except json.JSONDecodeError as e:
msg = f"Error: The file at '{json_path}' is not a valid JSON file: {e}"
logger.error(msg)
raise ValueError(msg)
except Exception as e:
logger.error(f"An unexpected error occurred: {e}")
return None
@@ -125,6 +127,8 @@ def get_models_full(primary_stem=None, model_t=None, file_t=None, json_path=None
# Allow for multiple values in the filters
if isinstance(primary_stem, str):
primary_stem = {primary_stem}
elif isinstance(primary_stem, list):
primary_stem = set(primary_stem)
if isinstance(model_t, str):
model_t = {model_t}
if isinstance(file_t, str):
@@ -153,8 +157,14 @@ def get_models_full(primary_stem=None, model_t=None, file_t=None, json_path=None
continue
# Check the stem
try:
if primary_stem is not None and d['primary_stem'] not in primary_stem:
continue
if primary_stem is not None:
mps = d['primary_stem']
if isinstance(mps, str):
if mps not in primary_stem:
continue
else: # A list
if not (set(mps) & primary_stem):
continue
except KeyError:
logger.error(f"Missing `primary_stem` for {name}")
continue
@@ -201,6 +211,7 @@ def get_models_full(primary_stem=None, model_t=None, file_t=None, json_path=None
on_disk_as_down.append(ICON_REMOTE + filtered_name)
else:
to_down.append(ICON_REMOTE + filtered_name)
d['hash'] = hash
found[filtered_name] = d
found_hashes[hash] = d
if file_name is not None:
@@ -228,14 +239,19 @@ def get_download_url(data):
return None
def cli_add_db(parser, default_json_file=None):
default_json_file = default_json_file or get_db_filename()
parser.add_argument('--json_file', type=str, default=default_json_file,
help="Path to the models database JSON file.")
def cli_add_models_and_db(parser):
# Compute the models dir assuming the script is run from a clone of the repo
default_json_file = get_db_filename()
parser.add_argument('--models_dir', type=str, default=os.path.dirname(default_json_file),
help="Path to the directory containing model files.")
parser.add_argument('--json_file', type=str, default=default_json_file,
help="Path to the models database JSON file.")
cli_add_db(parser, default_json_file=default_json_file)
def download_model(data, models_dir):
@@ -305,13 +321,42 @@ class ModelsDB(object):
def __init__(self, models_dir: str, json_path: str = None):
super().__init__()
self.models_dir = models_dir
self.json_path = json_path
self.json_path = get_db_filename(json_path)
self.refresh()
def refresh(self):
self.downloaded = hash_dir(self.models_dir)
self.models = load_known_models(self.json_path)
def remove(self, data):
if data is None:
return
del self.models[data["hash"]]
def add(self, hash, data):
self.models[hash] = data
def save(self):
# Remove run-time information
for k, v in self.models.items():
try:
del v["hash"]
except KeyError:
pass
try:
del v["indicator"]
except KeyError:
pass
try:
del v["filtered_name"]
except KeyError:
pass
try:
del v["model_path"]
except KeyError:
pass
save_known_models(self.models, self.json_path)
def get_filtered(self, primary_stem=None, model_t=None, file_t=None, default=None, repeat_dl=False):
return FilteredModels(primary_stem=primary_stem, model_t=model_t, file_t=file_t, json_path=self.json_path,
downloaded=self.downloaded, default=default, repeat_dl=repeat_dl)
source/inference/demixer.py
+198 -6
View File
@@ -5,6 +5,7 @@
#
# Wrappers for the model and inference
import logging
import math
import torch
# ComfyUI imports
try:
@@ -16,22 +17,42 @@ except Exception:
from .stft import stft_chunk_process, stft_get_chunks
from ..db.load_model import load_model
from ..utils.misc import NODES_NAME
from ..utils.torch import model_to_target, get_offload_device
from .demucs_api import apply_model, BagOfModels
logger = logging.getLogger(f"{NODES_NAME}.demixer")
SAMPLE_RATE = 44100
class DemixerGeneric(object):
def __init__(self, d, device, models_dir):
super().__init__()
self.d = d
self.model_run = load_model(d, device, models_dir)
self.device = device
def set_device(self, device):
if device == self.device:
return
self.device = device
if hasattr(self.model_run, "target_device"):
self.model_run = device
# ############################################################################################################################
# MDX-Net
# ############################################################################################################################
def show_inference_parameters(d):
logger.debug("Using inference parameters:")
logger.debug(f" Frequency Bins (n_fft/2): {d['mdx_n_fft_scale_set']//2}")
logger.debug(f" Amplitude Compensation: {d['compensate']}")
class DemixerMDX(object):
class DemixerMDX(DemixerGeneric):
def __init__(self, d, device, models_dir):
self.d = d
self.model_run = load_model(d, device, models_dir)
self.device = device
super().__init__(d, device, models_dir)
show_inference_parameters(d)
self.sr = SAMPLE_RATE
self.ch = 2
@@ -99,6 +120,177 @@ class DemixerMDX(object):
raise e
# ############################################################################################################################
# Demucs
# ############################################################################################################################
def get_steps_for_demucs(model, wav, segment, shifts, overlap):
segment = segment or model.segment
segment_length = int(model.samplerate * segment)
stride = int((1 - overlap) * segment_length)
return math.ceil(wav.shape[-1] / stride) * (shifts + 1)
class DemixerDemucs(DemixerGeneric):
def __init__(self, d, device, models_dir):
super().__init__(d, device, models_dir)
self.sr = self.model_run.samplerate
# Demucs code will move the model to and from the device
# The advantage is that it will be do it for the sub_models
# So here we keep it offloaded and tell Demucs code to do the work
self.model_run.target_device = get_offload_device()
def get_steps(self, wav, segment, shifts, overlap):
""" Tries to figure out how much steps we will need for inference """
if isinstance(self.model_run, BagOfModels):
total = 0
for sub_model in self.model_run.models:
steps = get_steps_for_demucs(sub_model, wav, segment, shifts, overlap)
total += steps
logger.debug(f"- Steps {steps}")
logger.debug(f"Total steps {total}")
return total
steps = get_steps_for_demucs(self.model_run, wav, segment, shifts, overlap)
logger.debug(f"Steps {steps}")
return steps
def demucs_callback(self, v):
self.comfy_progress_bar.update(1)
def __call__(self, waveform_tensor, segment=None, shifts=0, overlap=0.25):
try:
# --- 1. Normalize input shape to handle both batched and non-batched data ---
if waveform_tensor.ndim == 2:
# Input is [C, samples], add a batch dimension to make it [1, C, samples]
logger.debug("Input is not batched. Adding a temporary batch dimension.")
waveform_tensor = waveform_tensor.unsqueeze(0)
input_was_batched = False
elif waveform_tensor.ndim == 3:
# Input is already batched [B, C, samples]
input_was_batched = True
else:
raise ValueError(f"Unsupported waveform shape: {waveform_tensor.shape}. Expected 2 or 3 dimensions.")
batch_size = waveform_tensor.shape[0]
logger.info("🎛️ Performing demix...")
model = self.model_run
input_tensor_on_device = waveform_tensor.to(self.device)
# Determine the segment size
# 1. User selection
# 2. Value in the config
# 3. Auto: from each model (when None)
forced_segment = segment
if forced_segment is None:
forced_segment = self.model_run.config_segment
if forced_segment:
logger.debug(f"Using model provided segment size {forced_segment} s")
else:
logger.debug("Using default segment size")
forced_segment = None
else:
logger.debug(f"Using user provided segment size {forced_segment} s")
if with_comfy:
comfy_progress_bar = comfy.utils.ProgressBar(self.get_steps(waveform_tensor, forced_segment, shifts, overlap))
with model_to_target(model):
separated_tensors = apply_model(
model,
input_tensor_on_device,
device=self.device,
segment=forced_segment,
shifts=shifts + 1, # Shifts 0 is disabled, 1 is just one pass, 2 is 2 passes
overlap=overlap,
split=True, # Enable chunking
progress=True, # Show a progress bar in the console
callback=lambda x: (x.get('state') == 'end') and comfy_progress_bar.update(1) if with_comfy else None,
)
# Move the final result tensor back to the CPU before creating the output dicts.
# This is good practice to free up VRAM for subsequent nodes.
separated_tensors = separated_tensors.cpu()
assert batch_size == separated_tensors.shape[0]
# The output is [batch, sources, channels, samples].
# But we will separate the stems and return a batch for each stem, so we need
# [sources, batch, channels, samples].
separated_sources = separated_tensors.permute(1, 0, 2, 3)
# The model object tells us the names of the stems it produced
model_stems = model.sources.copy()
# UVR model is a 2 stems model with vocals and non_vocals, map it gracefully
if model_stems[1] == 'non_vocals':
model_stems[1] = 'other'
logger.debug(f"Model produced {len(model_stems)} stems: {model_stems}")
# Create a dictionary mapping the stem name to the audio tensor
output_map = dict(zip(model_stems, separated_sources))
# --- Silent audio for missing stems
reference_tensor = None
# Find the first valid tensor from our output to use as a shape reference
for tensor in output_map.values():
if tensor is not None and isinstance(tensor, torch.Tensor):
reference_tensor = tensor
break
# Determine the shape and sample rate for our silent audio fallback
if reference_tensor is not None:
# If we have a successful stem, use its properties
ref_batch_size, ref_channels, ref_samples = reference_tensor.shape
else:
# EDGE CASE: All stems failed. Fall back to the input audio's properties.
logger.warning("All stems failed. Using input audio shape for silence.")
# input_audio['waveform'] has shape [batch, channels, samples]
ref_batch_size, ref_channels, ref_samples = waveform_tensor.shape
# ---
# --- Gracefully create the 6 outputs ---
# Iterate through our fixed RETURN_NAMES and get the corresponding tensor.
# If a stem name doesn't exist in our output_map (e.g., 'guitar' for a 4-stem model),
# the .get() method will return None, which ComfyUI handles correctly.
final_outputs = []
for stem_name in ("vocals", "drums", "bass", "other", "guitar", "piano"):
output_tensor = output_map.get(stem_name, None)
# If the stem was produced, wrap it back into an AUDIO dict.
# If not, append None. ComfyUI handles None outputs correctly.
if output_tensor is not None:
if not input_was_batched:
output_tensor = output_tensor.squeeze(0)
output_dict = {
"waveform": output_tensor,
"sample_rate": model.samplerate,
"stem": stem_name.capitalize()
}
final_outputs.append(output_dict)
else:
logger.debug(f"Stem '{stem_name}' not produced by model. Returning silence.")
# Create a silent tensor with the correct shape, device, and dtype
silent_waveform = torch.zeros(ref_batch_size, ref_channels, ref_samples, dtype=torch.float32)
# Wrap the silent tensor in the ComfyUI AUDIO dictionary format
silent_dict = {
"waveform": silent_waveform,
"sample_rate": model.samplerate,
"stem": stem_name.capitalize()
}
final_outputs.append(silent_dict)
return tuple(final_outputs)
except Exception as e:
logger.error(f"Error during separation: {str(e)}")
raise e
def get_demixer(d, device, models_dir):
# Currently just MDX
return DemixerMDX(d, device, models_dir)
model_t = d['model_t'].lower()
if model_t == "mdx":
return DemixerMDX(d, device, models_dir)
elif model_t == "demucs":
return DemixerDemucs(d, device, models_dir)
msg = f"Unknown model type `{model_t}`"
logger.error(msg)
raise ValueError(msg)
source/inference/get_model.py
+103 -7
View File
@@ -4,18 +4,114 @@
# Project: ComfyUI-AudioSeparation
#
# Helper to get a model from the correct class
import importlib
import json
import logging
from safetensors import safe_open
from .MDX_Net import MDX_Net
from ..utils.misc import NODES_NAME
from ..utils.misc import NODES_NAME, json_object_hook
# Demucs class imports
from .demucs_api import BagOfModels
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_metadata(file_path):
""" Read the metadata from a safetensors file """
logger.debug(f"Reading metadata from {file_path}")
metadata = {}
with safe_open(file_path, framework="pt", device="cpu") as f:
metadata = f.metadata()
if not metadata:
raise ValueError(f"Could not read metadata from safetensors file: {file_path}")
return metadata
def get_hyperparameter(metadata, parameter, as_type, default=None, warn_diff=True):
value = metadata.get(parameter)
if value is None:
if default is None:
raise ValueError(f"Missing `{parameter}` hyperparameter")
return default
if as_type == "int":
value = int(value)
if warn_diff and value != default:
logger.warning(f"Hyperparameter mismatch: database = {default}, metadata = {value}")
return value
def get_mdx_model(d):
""" Create an MDX_Net object with the specified parameters """
# Check the file is consistent we our data base
metadata = get_metadata(d['model_path'])
dim_f = get_hyperparameter(metadata, 'mdx_dim_f_set', "int", d['mdx_dim_f_set'])
channels = get_hyperparameter(metadata, 'channels', "int", d['channels'])
stages = get_hyperparameter(metadata, 'stages', "int", d['stages'])
# Create a class with this parameters
return MDX_Net(dim_f=dim_f, ch=channels, num_stages=stages)
def get_demucs_model(d):
""" Create a Demucs, HDemucs (Hybrid) or HTDemucs (Hybrid Transformer) object.
All metadata comes from the safetensors """
file_path = d['model_path']
# 1. First, open the file safely to read only the metadata header.
metadata = get_metadata(file_path)
is_bag = json.loads(metadata.get('is_bag_of_models', 'false'))
signatures = json.loads(metadata['signatures'])
sub_models = []
for sig in set(signatures): # Use set to only instantiate each unique architecture once
model_meta_str = metadata.get(sig)
if not model_meta_str:
raise ValueError(f"Metadata for signature '{sig}' not found in safetensors file.")
# Use the object_hook here to reconstruct Fraction objects automatically
model_meta = json.loads(model_meta_str, object_hook=json_object_hook)
class_module, class_name = model_meta['class_module'], model_meta['class_name']
args, kwargs = model_meta['args'], model_meta['kwargs']
logger.debug(f" - Reconstructing architecture for '{sig}': {class_module}.{class_name}")
assert '.' not in class_name, "Security check failed, won't import a file outside my directory"
# Import the class from a module in this dir with the same name as the class
local_class_module = '.'.join(__name__.split('.')[:-1]) + "." + class_name
logger.debug(f" - Redirecting class: {class_module} -> {local_class_module}")
module = importlib.import_module(local_class_module)
klass = getattr(module, class_name)
sub_models.append({'sig': sig, 'model': klass(*args, **kwargs)})
# Create a mapping from signature to model instance
model_map = {m['sig']: m['model'] for m in sub_models}
# Re-order the models to match the YAML's signature list
ordered_models = [model_map[sig] for sig in signatures]
segment = float(metadata.get('segment', '0'))
if not is_bag:
final_model = ordered_models[0]
final_model.signatures = signatures
else:
logger.debug("Rebuilding BagOfModels container...")
weights = json.loads(metadata.get('weights', 'null'))
final_model = BagOfModels(ordered_models, weights=weights, segment=segment)
final_model.signatures = signatures
final_model.config_segment = segment
return final_model
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'])
if model_t == "mdx":
return get_mdx_model(d)
elif model_t == "demucs":
return get_demucs_model(d)
msg = f"Unknown model type `{model_t}`"
logger.error(msg)
raise ValueError(msg)
source/inference/stft.py
-1
View File
@@ -126,7 +126,6 @@ def stft_chunk_process(waveform, d, model_run, device, segment_size=256, hop_len
total_chunks = 1 + (mixture.shape[1] - chunk_size + step - 1) // step
logger.info(f"⚙️ Processing {total_chunks} chunks...")
model_run.target_device = device
with model_to_target(model_run):
for i in tqdm(range(0, mixture.shape[1] - chunk_size + 1, step)):