Meisoftcoltd
4bfe6b3767
- Add ComfyUI "tooltip" property to all inputs in the `DemucsAudioSeparator` node - Tooltips are provided in Spanish to explain functionality like `shifts`, `overlap`, `split`, etc. - Did not add `bfloat16` to options as it's unsupported by cuFFT and would downgrade to `float32` anyways.
303 lines
13 KiB
Python
303 lines
13 KiB
Python
import torch
|
|
import torchaudio
|
|
import os
|
|
import json
|
|
import numpy as np
|
|
import inspect
|
|
from pathlib import Path
|
|
|
|
# Try to import ComfyUI utilities
|
|
try:
|
|
import folder_paths
|
|
except ImportError:
|
|
# Fallback for development/testing outside ComfyUI
|
|
class folder_paths:
|
|
models_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "models")
|
|
@staticmethod
|
|
def get_output_directory():
|
|
return os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "output")
|
|
|
|
try:
|
|
import comfy.utils
|
|
HAS_COMFY = True
|
|
except ImportError:
|
|
HAS_COMFY = False
|
|
|
|
from demucs.pretrained import get_model
|
|
from demucs.apply import apply_model
|
|
|
|
class LoadAudioDirectory:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"directory_path": ("STRING", {"default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("AUDIO",)
|
|
RETURN_NAMES = ("audio",)
|
|
OUTPUT_IS_LIST = (True,)
|
|
FUNCTION = "load_directory"
|
|
CATEGORY = "🎵 Demucs-AudioSeparator ⚡"
|
|
|
|
def load_directory(self, directory_path):
|
|
if not os.path.isdir(directory_path):
|
|
raise ValueError(f"⚡ Directory not found: {directory_path}")
|
|
|
|
supported_exts = {".wav", ".mp3", ".flac", ".ogg", ".m4a"}
|
|
|
|
audio_files = []
|
|
for file in sorted(os.listdir(directory_path)):
|
|
ext = os.path.splitext(file)[1].lower()
|
|
if ext in supported_exts:
|
|
audio_files.append(os.path.join(directory_path, file))
|
|
|
|
if not audio_files:
|
|
raise ValueError(f"⚡ No supported audio files found in directory: {directory_path}")
|
|
|
|
print(f"⚡ Loading {len(audio_files)} audio files from {directory_path}...")
|
|
|
|
audios = []
|
|
for file_path in audio_files:
|
|
try:
|
|
try:
|
|
# Attempt 1: Force 'soundfile' backend to bypass the default torchcodec
|
|
waveform, sample_rate = torchaudio.load(file_path, backend='soundfile')
|
|
except Exception as e:
|
|
print(f"⚡ Warning: torchaudio failed (likely torchcodec issue). Falling back to librosa... Error: {e}")
|
|
import librosa
|
|
# Attempt 2: Use librosa which reads MP3s without relying on torchcodec
|
|
audio_np, sample_rate = librosa.load(file_path, sr=None, mono=False)
|
|
|
|
# Ensure shape is [channels, frames]
|
|
if audio_np.ndim == 1:
|
|
audio_np = np.expand_dims(audio_np, axis=0)
|
|
|
|
waveform = torch.from_numpy(audio_np)
|
|
|
|
# Ensure float32 precision
|
|
waveform = waveform.to(torch.float32)
|
|
|
|
# Add batch dimension to match standard audio format [batch, channels, samples]
|
|
if waveform.dim() == 2:
|
|
waveform = waveform.unsqueeze(0)
|
|
|
|
audios.append({"waveform": waveform, "sample_rate": sample_rate})
|
|
print(f"⚡ Successfully loaded {os.path.basename(file_path)}")
|
|
except Exception as e:
|
|
print(f"⚡ Error loading {file_path}: {e}")
|
|
|
|
if not audios:
|
|
raise ValueError(f"⚡ Failed to load any audio files from {directory_path}")
|
|
|
|
return (audios,)
|
|
|
|
|
|
class DemucsAudioSeparator:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio": ("AUDIO", {"tooltip": "El clip de audio a procesar."}),
|
|
"model": (["htdemucs", "htdemucs_ft", "htdemucs_6s", "hdemucs_mmi", "mdx", "mdx_extra", "mdx_q", "mdx_extra_q", "mdxc", "mdxc_fb_ft"], {"tooltip": "El modelo de Demucs a utilizar para la separación."}),
|
|
"device": (["cuda", "cpu"], {"tooltip": "El dispositivo a usar para la inferencia. 'cuda' para GPU, 'cpu' para procesador."}),
|
|
"precision": (["float32", "float16"], {"default": "float32", "tooltip": "Precisión de los cálculos. 'float32' es más seguro, 'float16' es más rápido pero menos estable."}),
|
|
"shifts": ("INT", {"default": 1, "min": 1, "max": 10, "tooltip": "Número de predicciones aleatorias (mayor valor = mayor calidad, pero más lento)."}),
|
|
"overlap": ("FLOAT", {"default": 0.25, "min": 0.1, "max": 0.9, "step": 0.05, "tooltip": "Solapamiento entre ventanas de procesamiento. Por defecto: 0.25."}),
|
|
"split": ("BOOLEAN", {"default": True, "tooltip": "Divide el audio en segmentos para ahorrar memoria VRAM. Imprescindible para audios largos."}),
|
|
"vocals": ("BOOLEAN", {"default": True, "tooltip": "Extraer la pista de voces."}),
|
|
"drums": ("BOOLEAN", {"default": True, "tooltip": "Extraer la pista de batería."}),
|
|
"bass": ("BOOLEAN", {"default": True, "tooltip": "Extraer la pista de bajo."}),
|
|
"other": ("BOOLEAN", {"default": True, "tooltip": "Extraer otros instrumentos."}),
|
|
"guitar": ("BOOLEAN", {"default": False, "tooltip": "Extraer la pista de guitarra (solo modelos de 6 pistas)."}),
|
|
"piano": ("BOOLEAN", {"default": False, "tooltip": "Extraer la pista de piano (solo modelos de 6 pistas)."}),
|
|
"instrumental": ("BOOLEAN", {"default": True, "tooltip": "Extraer la pista instrumental (audio original sin voces)."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("AUDIO", "AUDIO", "AUDIO", "AUDIO", "AUDIO", "AUDIO", "AUDIO", "JSON")
|
|
RETURN_NAMES = ("vocals", "drums", "bass", "other", "guitar", "piano", "instrumental", "metadata")
|
|
FUNCTION = "separate"
|
|
CATEGORY = "🎵 Demucs-AudioSeparator ⚡"
|
|
|
|
def separate(self, audio, model, device, precision, shifts, overlap, split, vocals, drums, bass, other, guitar, piano, instrumental):
|
|
model_name = model
|
|
|
|
# Legacy support and safety check for bfloat16
|
|
if precision == "bfloat16":
|
|
print("⚠️ [Demucs Pro] BFloat16 not supported for FFT. Upgrading to Float32 for inference.")
|
|
precision = "float32"
|
|
|
|
if precision == "float16":
|
|
dtype = torch.float16
|
|
else:
|
|
dtype = torch.float32
|
|
|
|
# Configure model path
|
|
demucs_models_path = os.path.join(folder_paths.models_dir, "demucs")
|
|
if not os.path.exists(demucs_models_path):
|
|
os.makedirs(demucs_models_path, exist_ok=True)
|
|
|
|
torch.hub.set_dir(demucs_models_path)
|
|
|
|
# Determine device
|
|
if device == "cuda" and not torch.cuda.is_available():
|
|
print("⚡ CUDA not available, falling back to CPU")
|
|
device = "cpu"
|
|
|
|
device_obj = torch.device(device)
|
|
|
|
# Load model
|
|
print(f"⚡ Loading Demucs model: {model}...")
|
|
try:
|
|
model_inst = get_model(model)
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to load Demucs model {model}: {str(e)}")
|
|
|
|
model_inst.to(device_obj)
|
|
model_inst.eval()
|
|
|
|
# Prepare audio
|
|
waveform = audio["waveform"]
|
|
sr = audio["sample_rate"]
|
|
|
|
# Ensure waveform is [batch, channels, samples]
|
|
if waveform.dim() == 2:
|
|
waveform = waveform.unsqueeze(0)
|
|
|
|
# --- FIX 1: CHANNEL MANAGEMENT (MONO/STEREO/SURROUND) ---
|
|
# Demucs strictly requires 2 channels (stereo)
|
|
channels = waveform.shape[1]
|
|
|
|
if channels == 1:
|
|
print("⚡ Detected mono audio, converting to stereo for Demucs...")
|
|
waveform = waveform.repeat(1, 2, 1)
|
|
elif channels > 2:
|
|
print(f"⚡ Detected {channels} channels, clipping to stereo for Demucs...")
|
|
waveform = waveform[:, :2, :]
|
|
|
|
# --- FIX 2: PRECISION SAFETY ---
|
|
# Move to device BUT keep in float32 for resampling
|
|
# (Prevents "Input type HalfTensor vs Weight type FloatTensor" error)
|
|
waveform = waveform.to(device_obj).to(torch.float32)
|
|
|
|
# Resample if necessary
|
|
if sr != model_inst.samplerate:
|
|
print(f"⚡ Resampling audio from {sr} to {model_inst.samplerate}...")
|
|
resampler = torchaudio.transforms.Resample(sr, model_inst.samplerate).to(device_obj)
|
|
waveform = resampler(waveform)
|
|
|
|
# Convert to target precision (Note: apply_model might force float32 later for stability)
|
|
waveform = waveform.to(dtype)
|
|
|
|
# Progress bar integration
|
|
pbar = None
|
|
if HAS_COMFY:
|
|
pbar = comfy.utils.ProgressBar(100)
|
|
|
|
def progress_callback(info):
|
|
if pbar:
|
|
# Support various keys from different Demucs versions
|
|
# e.g., 'progress', 'shift', 'total', 'shift_idx'
|
|
total = info.get('total') or info.get('shifts') or shifts
|
|
current = info.get('shift') or info.get('shift_idx') or info.get('progress')
|
|
|
|
if isinstance(current, float) and current <= 1.0:
|
|
pbar.update_absolute(int(current * 100))
|
|
elif total and current is not None:
|
|
# current might be 0-indexed
|
|
p_val = min(100, int((current + 1) / total * 100))
|
|
pbar.update_absolute(p_val)
|
|
|
|
# Apply model
|
|
print(f"⚡ Separating audio with {model} (shifts={shifts}, overlap={overlap}, split={split})...")
|
|
try:
|
|
with torch.no_grad():
|
|
# Check if the installed version of apply_model supports a callback
|
|
apply_kwargs = {
|
|
"shifts": shifts,
|
|
"split": split,
|
|
"overlap": overlap,
|
|
"progress": False,
|
|
"device": device_obj,
|
|
}
|
|
|
|
sig = inspect.signature(apply_model)
|
|
if 'callback' in sig.parameters:
|
|
apply_kwargs["callback"] = progress_callback
|
|
|
|
# Force float32 for inference to avoid cuFFT/BFloat16 issues
|
|
out = apply_model(model_inst, waveform.to(torch.float32), **apply_kwargs)
|
|
except Exception as e:
|
|
raise RuntimeError(f"⚡ [Demucs Pro] Error during Demucs inference: {str(e)}")
|
|
|
|
# out shape: [batch, sources, channels, samples]
|
|
sources = model_inst.sources
|
|
results = {}
|
|
for i, source_name in enumerate(sources):
|
|
# Convert back to float32 for output and move to CPU
|
|
results[source_name] = {
|
|
"waveform": out[:, i, :, :].to(torch.float32).cpu(),
|
|
"sample_rate": model_inst.samplerate
|
|
}
|
|
|
|
# Helper to get stem or zeroed audio based on user selection
|
|
def get_stem(name, enabled):
|
|
if enabled and name in results:
|
|
return results[name]
|
|
else:
|
|
# Return zeroed audio with same length and batch size if stem is disabled or unavailable
|
|
batch_size = out.shape[0]
|
|
channels = out.shape[2]
|
|
samples = out.shape[3]
|
|
return {
|
|
"waveform": torch.zeros((batch_size, channels, samples), dtype=torch.float32),
|
|
"sample_rate": model_inst.samplerate
|
|
}
|
|
|
|
# Map results to discrete outputs, respecting the boolean flags
|
|
out_vocals = get_stem("vocals", vocals)
|
|
out_drums = get_stem("drums", drums)
|
|
out_bass = get_stem("bass", bass)
|
|
out_other = get_stem("other", other)
|
|
out_guitar = get_stem("guitar", guitar)
|
|
out_piano = get_stem("piano", piano)
|
|
|
|
if instrumental:
|
|
if "vocals" in results:
|
|
# Subtract vocals from original resampled waveform
|
|
inst_waveform = waveform.cpu() - results["vocals"]["waveform"]
|
|
else:
|
|
# If the model doesn't output vocals, just return the original audio
|
|
inst_waveform = waveform.cpu()
|
|
|
|
out_instrumental = {
|
|
"waveform": inst_waveform.to(torch.float32),
|
|
"sample_rate": model_inst.samplerate
|
|
}
|
|
else:
|
|
batch_size = out.shape[0]
|
|
channels = out.shape[2]
|
|
samples = out.shape[3]
|
|
out_instrumental = {
|
|
"waveform": torch.zeros((batch_size, channels, samples), dtype=torch.float32),
|
|
"sample_rate": model_inst.samplerate
|
|
}
|
|
|
|
metadata = {
|
|
"model": model_name,
|
|
"device": str(device_obj),
|
|
"precision": precision,
|
|
"shifts": shifts,
|
|
"overlap": overlap,
|
|
"split": split,
|
|
"available_sources": sources,
|
|
"input_sample_rate": sr,
|
|
"output_sample_rate": model_inst.samplerate,
|
|
"status": "success"
|
|
}
|
|
|
|
print(f"⚡ [Demucs Pro] Separation completed successfully.")
|
|
return (out_vocals, out_drums, out_bass, out_other, out_guitar, out_piano, out_instrumental, metadata)
|