134 lines
4.5 KiB
Python
134 lines
4.5 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: UTF-8 -*-
|
|
import os
|
|
import random
|
|
import torch
|
|
import numpy as np
|
|
import logging
|
|
import folder_paths
|
|
from pathlib import Path
|
|
import torchaudio
|
|
|
|
from .demucs.api import Separator,save_audio
|
|
from .demucs.audio import pre_audio
|
|
node_cur_path = os.path.dirname(os.path.abspath(__file__))
|
|
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
|
|
|
|
|
class Demucs_Loader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"d_model": ("STRING", {"default": "htdemucs"}),
|
|
"overlap": ("FLOAT", {"default": 0.25, "min": 0.01, "max": 1.0, "step": 0.01,}),
|
|
"shifts": (
|
|
"INT", {"default": 1, "min": 1, "max": 4096, "step": 1, "display": "number"}),
|
|
"split": ("BOOLEAN", {"default": True},),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("Demucs_MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "main_"
|
|
CATEGORY = "Demucs"
|
|
|
|
def main_(self, d_model,overlap,shifts,split):
|
|
repo=None
|
|
model = Separator(
|
|
model=d_model,
|
|
repo=repo,
|
|
device=device,
|
|
shifts=shifts,
|
|
overlap=overlap,
|
|
split=split,
|
|
segment=None,
|
|
jobs=0,
|
|
callback=print
|
|
)
|
|
# out = args.out / args.name
|
|
# out.mkdir(parents=True, exist_ok=True)
|
|
logging.info("loading checkpoint done.")
|
|
|
|
return (model,)
|
|
|
|
|
|
class Demucs_Sampler:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("Demucs_MODEL",),
|
|
"audio": ("AUDIO",),
|
|
"ext": (["mp3","flac","wav"],),
|
|
"bits_per_sample": ([24, 16,],),
|
|
"as_float": (["float32", "float16", ],),
|
|
"clip_mode": (["rescale", "clamp", "tanh","none"],),
|
|
"mp3_bitrate": (
|
|
"INT", {"default": 320, "min": 1, "max": 4096, "step": 1, "display": "number"}),
|
|
"audio_save": ("BOOLEAN", {"default": True},),
|
|
"preset": ([2,3,4,5,6,7],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("AUDIO","AUDIO","AUDIO","AUDIO","AUDIO",)
|
|
RETURN_NAMES = ("mixed","bass","drums","other","vocals",)
|
|
FUNCTION = "main"
|
|
CATEGORY = "Demucs"
|
|
|
|
def main(self, model, audio, ext, bits_per_sample, as_float, clip_mode, mp3_bitrate,audio_save,preset):
|
|
|
|
#pre input audio
|
|
audio_file_prefix = ''.join(random.choice("0123456789abcdefg") for _ in range(6))
|
|
file = os.path.join(folder_paths.get_input_directory(), f"audio_{audio_file_prefix}.{ext}")
|
|
torchaudio.save(file, audio["waveform"].squeeze(0), audio["sample_rate"])
|
|
|
|
out=Path(folder_paths.get_output_directory())
|
|
|
|
separated = model.separate_audio_file(file)[1]
|
|
kwargs = {
|
|
"samplerate": model.samplerate,
|
|
"bitrate": mp3_bitrate,
|
|
"clip": clip_mode,
|
|
"as_float": as_float,
|
|
"bits_per_sample": bits_per_sample,
|
|
}
|
|
samplerate=44100
|
|
audio_list=[]
|
|
for stem, source in separated.items():
|
|
audio_dict=pre_audio(source,ext,samplerate,clip_mode)
|
|
audio_list.append(audio_dict)
|
|
if audio_save:
|
|
stem = out / "{track}/{stem}.{ext}".format(
|
|
track=Path(file).name.rsplit(".", 1)[0],
|
|
trackext=Path(file).name.rsplit(".", 1)[-1],
|
|
stem=stem,
|
|
ext=ext,
|
|
)
|
|
stem.parent.mkdir(parents=True, exist_ok=True)
|
|
save_audio(source, str(stem), **kwargs)
|
|
output_ = os.path.join(folder_paths.get_output_directory(), f"audio_{audio_file_prefix}")
|
|
logging.info(f"save audio in {output_} dir...")
|
|
|
|
mixed_wav=audio_list[0]["waveform"]+audio_list[1]["waveform"]+audio_list[2]["waveform"]
|
|
mixed={"waveform": mixed_wav, "sample_rate": samplerate}
|
|
return (mixed,audio_list[0],audio_list[1],audio_list[2],audio_list[3],)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Demucs_Loader": Demucs_Loader,
|
|
"Demucs_Sampler": Demucs_Sampler,
|
|
}
|
|
|
|
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Demucs_Loader": "Demucs_Loader",
|
|
"Demucs_Sampler": "Demucs_Sampler",
|
|
}
|