130 lines
5.5 KiB
Python
130 lines
5.5 KiB
Python
import os
|
|
import folder_paths
|
|
import torch
|
|
import torchaudio
|
|
import tempfile
|
|
import shutil
|
|
from audio_separator.separator import Separator
|
|
|
|
# Определяем путь к моделям: ComfyUI/models/UVR
|
|
COMFY_MODELS_DIR = folder_paths.models_dir
|
|
UVR_MODELS_DIR = os.path.join(COMFY_MODELS_DIR, "UVR")
|
|
|
|
# Если папки нет - создаем
|
|
if not os.path.exists(UVR_MODELS_DIR):
|
|
try:
|
|
os.makedirs(UVR_MODELS_DIR)
|
|
except Exception:
|
|
pass
|
|
|
|
class AudioSeparatorNode:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
# Сканируем папку models/UVR на наличие моделей
|
|
model_files = []
|
|
if os.path.exists(UVR_MODELS_DIR):
|
|
model_files = [f for f in os.listdir(UVR_MODELS_DIR) if f.endswith(('.pth', '.onnx', '.bin'))]
|
|
|
|
# Если папка пустая, добавляем дефолтные имена, чтобы список не был пустым
|
|
if not model_files:
|
|
model_files = ["1_HP-UVR.pth", "UVR-DeNoise.pth", "UVR-MDX-NET-Inst_HQ_3.onnx"]
|
|
|
|
return {
|
|
"required": {
|
|
"audio": ("AUDIO",),
|
|
"model_filename": (model_files, {"default": model_files[0] if model_files else "1_HP-UVR.pth"}),
|
|
"window_size": ("INT", {"default": 512, "min": 320, "max": 1024, "step": 64, "tooltip": "VR Model Window Size. 512 is High Quality."}),
|
|
"aggression": ("INT", {"default": 10, "min": 1, "max": 50, "step": 1, "tooltip": "VR Model Aggression. Higher = stronger removal."}),
|
|
"batch_size": ("INT", {"default": 4, "min": 1, "max": 16, "step": 1, "tooltip": "Batch size for GPU processing."}),
|
|
},
|
|
"optional": {
|
|
"normalization": ("FLOAT", {"default": 0.9, "min": 0.1, "max": 1.0, "step": 0.1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("AUDIO", "AUDIO")
|
|
RETURN_NAMES = ("Vocals (Clean)", "Instrumental (Noise)")
|
|
FUNCTION = "separate_audio"
|
|
CATEGORY = "Audio/Processing"
|
|
|
|
def separate_audio(self, audio, model_filename, window_size, aggression, batch_size, normalization):
|
|
# 1. Подготовка аудио из Comfy
|
|
waveform = audio['waveform']
|
|
# Берем sample_rate входного файла только для сохранения времянки
|
|
input_sample_rate = audio['sample_rate']
|
|
|
|
# Если батч > 1, берем только первый элемент
|
|
if waveform.dim() == 3:
|
|
waveform = waveform[0]
|
|
|
|
# Создаем временную папку
|
|
temp_dir = tempfile.mkdtemp()
|
|
input_path = os.path.join(temp_dir, "input_temp.wav")
|
|
|
|
# Сохраняем временный файл
|
|
torchaudio.save(input_path, waveform, input_sample_rate)
|
|
|
|
# 2. Инициализация Сепаратора
|
|
print(f"🎧 Audio Separator: Initializing with model '{model_filename}'...")
|
|
|
|
separator = Separator(
|
|
model_file_dir=UVR_MODELS_DIR, # Ищем модели в ComfyUI/models/UVR
|
|
output_dir=temp_dir, # Результат кладем во временную папку
|
|
output_format="WAV",
|
|
normalization_threshold=normalization,
|
|
vr_params={
|
|
"batch_size": batch_size,
|
|
"window_size": window_size,
|
|
"aggression": aggression,
|
|
},
|
|
mdx_params={
|
|
"batch_size": batch_size,
|
|
}
|
|
)
|
|
|
|
# Загружаем модель
|
|
try:
|
|
separator.load_model(model_filename=model_filename)
|
|
except Exception as e:
|
|
shutil.rmtree(temp_dir)
|
|
raise ValueError(f"Ошибка загрузки модели '{model_filename}'. Убедитесь, что она лежит в {UVR_MODELS_DIR}. \nОшибка: {e}")
|
|
|
|
# 3. Обработка
|
|
print(f"🎧 Audio Separator: Processing on GPU...")
|
|
output_files = separator.separate(input_path)
|
|
|
|
# Определяем, какой файл есть какой
|
|
vocals_path = os.path.join(temp_dir, output_files[0])
|
|
instr_path = os.path.join(temp_dir, output_files[1])
|
|
|
|
for f in output_files:
|
|
full_p = os.path.join(temp_dir, f)
|
|
if "(Vocals)" in f: vocals_path = full_p
|
|
if "(Instrumental)" in f or "(Noise)" in f: instr_path = full_p
|
|
|
|
# 4. Возврат в Comfy
|
|
try:
|
|
# ВАЖНО: Читаем Sample Rate из обработанного файла!
|
|
# Модели часто меняют частоту на 44100, и если вернуть старую - будет слоу-мо.
|
|
vocals_wave, vocals_sr = torchaudio.load(vocals_path)
|
|
instr_wave, instr_sr = torchaudio.load(instr_path)
|
|
except Exception as e:
|
|
shutil.rmtree(temp_dir)
|
|
raise RuntimeError(f"Ошибка чтения результата: {e}")
|
|
|
|
# Формируем формат (Batch, Channels, Samples)
|
|
vocals_out = {"waveform": vocals_wave.unsqueeze(0), "sample_rate": vocals_sr}
|
|
instr_out = {"waveform": instr_wave.unsqueeze(0), "sample_rate": instr_sr}
|
|
|
|
# Удаляем временные файлы
|
|
try:
|
|
shutil.rmtree(temp_dir)
|
|
except:
|
|
pass
|
|
|
|
print("🎧 Audio Separator: Done!")
|
|
|
|
return (vocals_out, instr_out) |