300 lines
13 KiB
Python
300 lines
13 KiB
Python
# Copyright (c) 2025 Salvador E. Tropea
|
|
# Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
|
|
# License: GPLv3
|
|
# Project: ComfyUI-AudioSeparation
|
|
#
|
|
# Wrappers for the model and inference
|
|
import logging
|
|
import math
|
|
import torch
|
|
# ComfyUI imports
|
|
try:
|
|
import comfy.utils
|
|
with_comfy = True
|
|
except Exception:
|
|
with_comfy = False
|
|
# Local imports
|
|
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(DemixerGeneric):
|
|
def __init__(self, d, device, models_dir):
|
|
super().__init__(d, device, models_dir)
|
|
show_inference_parameters(d)
|
|
self.sr = SAMPLE_RATE
|
|
self.ch = 2
|
|
|
|
def __call__(self, waveform, segments=None):
|
|
if segments is None:
|
|
# To make it compatible with Demucs
|
|
segments = 1
|
|
dim_t = (2 ** self.d['mdx_dim_t_set']) * segments
|
|
try:
|
|
# --- 1. Normalize input shape to handle both batched and non-batched data ---
|
|
if waveform.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 = waveform.unsqueeze(0)
|
|
input_was_batched = False
|
|
elif waveform.ndim == 3:
|
|
# Input is already batched [B, C, samples]
|
|
input_was_batched = True
|
|
else:
|
|
raise ValueError(f"Unsupported waveform shape: {waveform.shape}. Expected 2 or 3 dimensions.")
|
|
|
|
batch_size = waveform.shape[0]
|
|
logger.info("🎛️ Performing demix...")
|
|
|
|
# Lists to store the separated stems from each item in the batch
|
|
list_of_main_stems = []
|
|
list_of_complement_stems = []
|
|
|
|
# ComfyUI progress bar
|
|
progress_bar_ui = None
|
|
if with_comfy:
|
|
chunks = stft_get_chunks(waveform.shape[2], self.d['mdx_n_fft_scale_set'], segment_size=dim_t)
|
|
chunks *= batch_size
|
|
progress_bar_ui = comfy.utils.ProgressBar(chunks)
|
|
|
|
# --- 2. Iterate through the batch ---
|
|
for i, single_waveform in enumerate(waveform):
|
|
# single_waveform has shape [C, samples]
|
|
logger.debug(f"Processing item {i+1}/{batch_size}...")
|
|
|
|
# Process this single waveform
|
|
main_wav = stft_chunk_process(single_waveform, self.d, self.model_run, self.device, segment_size=dim_t,
|
|
progress_bar_ui=progress_bar_ui)
|
|
complement_wav = single_waveform - main_wav
|
|
|
|
# Add the results to our lists
|
|
list_of_main_stems.append(main_wav)
|
|
list_of_complement_stems.append(complement_wav)
|
|
|
|
# --- 3. Stack the results into single batch tensors ---
|
|
# torch.stack creates a new dimension (the batch dimension) from a list of tensors
|
|
stacked_main_stems = torch.stack(list_of_main_stems, dim=0)
|
|
stacked_complement_stems = torch.stack(list_of_complement_stems, dim=0)
|
|
# Both will now have shape [B, C, samples]
|
|
|
|
# --- 4. Denormalize output shape if original input was not batched ---
|
|
if not input_was_batched:
|
|
logger.debug("Squeezing batch dimension from output to match non-batched input.")
|
|
stacked_main_stems = stacked_main_stems.squeeze(0)
|
|
stacked_complement_stems = stacked_complement_stems.squeeze(0)
|
|
|
|
return [{'waveform': stacked_main_stems, 'sample_rate': SAMPLE_RATE, 'stem': self.d['primary_stem']},
|
|
{'waveform': stacked_complement_stems, 'sample_rate': SAMPLE_RATE, 'stem': 'Complement'}]
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error during separation: {str(e)}")
|
|
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):
|
|
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)
|