197 lines
8.2 KiB
Python
197 lines
8.2 KiB
Python
"""
|
|
╔════════════════════════════════════════╗
|
|
║ ║
|
|
║ Developed by Abdullah OZMANTAR ║
|
|
║ GitHub : github.com/abdozmantar ║
|
|
║ ║
|
|
╚════════════════════════════════════════╝
|
|
"""
|
|
|
|
import sys
|
|
|
|
from dora.log import fatal
|
|
import torch as th
|
|
from pathlib import Path
|
|
|
|
import random
|
|
import string
|
|
|
|
from .audio import save_audio
|
|
from .helper import Separator
|
|
from .pretrained import SOURCES
|
|
from .repo import ModelLoadingError
|
|
|
|
class MultiInput(str):
|
|
def __new__(cls, string, allowed_types="*"):
|
|
res = super().__new__(cls, string)
|
|
res.allowed_types=allowed_types
|
|
return res
|
|
def __ne__(self, other):
|
|
if self.allowed_types == "*" or other == "*":
|
|
return False
|
|
return other not in self.allowed_types
|
|
|
|
audioOrNone = MultiInput("AUDIO", ["AUDIO", "*"])
|
|
|
|
class DeepExtractV2Node:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"track": ("AUDIO", {"default": None, "multiline": True, "tooltip": "Input audio track to separate"}),
|
|
"out": ("STRING", {"default": "separated", "tooltip": "Folder where to put extracted track. This folder will be created inside the main ComfyUI directory."}),
|
|
"device": (["cuda", "cpu"], {"default": "cuda", "tooltip": "Device to use for inference"}),
|
|
"overlap": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "Overlap between audio chunks"}),
|
|
"shifts": ("INT", {"default": 1, "min": 1, "max": 20, "tooltip": "Number of random shifts for stabilization"}),
|
|
"segment": ("INT", {"default": 0, "min": 0, "max": 600, "tooltip": "Split size of each chunk (0 = no split)"}),
|
|
"bit_depth": (["default", "int24", "float32"], {"default": "default", "tooltip": "Audio output bit depth"}),
|
|
"clip_mode": (["rescale", "clamp", "none"], {"default": "rescale", "tooltip": "Clipping handling strategy"}),
|
|
"audio_format": (["wav", "flac", "mp3"], {"default": "wav", "tooltip": "Output audio file format"}),
|
|
"mp3_bitrate": ("INT", {"default": 320, "min": 64, "max": 320, "tooltip": "Bitrate for MP3 output"}),
|
|
"mp3_preset": ("INT", {"default": 2, "min": 2, "max": 7, "tooltip": "Encoder preset (2=best, 7=fast)"}),
|
|
"jobs": ("INT", {"default": 0, "min": 0, "max": 32, "tooltip": "Number of parallel jobs"}),
|
|
"split": ("BOOLEAN", {"default": True, "tooltip": "Split audio in chunks to save memory"}),
|
|
"other_method": (["none", "add", "minus"], {"default": "add","tooltip": 'Decide how to get "no_{STEM}". "none" will not save "no_{STEM}". ''"add" will add all the other stems. "minus" will use the original track minus the selected stem.'}),
|
|
"stem": ("STRING", { "default": None ,"tooltip": "Only separate audio into {STEM} and no_{STEM}."}),
|
|
"filename": ("STRING", {"default": "{track}/{stem}.{ext}", "tooltip": "Set the name of output file. Use '{track}', '{stem}', '{ext}' to use variables of track name without extension, track extension,stem name and default output file extension. Default is {track}/{stem}.{ext}"}),
|
|
}
|
|
}
|
|
|
|
OUTPUT_SOURCES = tuple(list(SOURCES) + ["derived_stem"])
|
|
|
|
RETURN_TYPES = tuple([audioOrNone] * len(OUTPUT_SOURCES))
|
|
RETURN_NAMES = tuple(OUTPUT_SOURCES)
|
|
|
|
FUNCTION = "separate"
|
|
CATEGORY = "DeepExtractV2"
|
|
|
|
|
|
|
|
def separate(self,
|
|
track,
|
|
out,
|
|
device,
|
|
overlap,
|
|
shifts,
|
|
segment,
|
|
bit_depth,
|
|
clip_mode,
|
|
audio_format,
|
|
mp3_bitrate,
|
|
mp3_preset,
|
|
jobs,
|
|
split,
|
|
other_method,
|
|
stem,
|
|
filename):
|
|
|
|
if(segment == 0):
|
|
segment = None
|
|
|
|
out = Path(out)
|
|
|
|
if len(track) == 0:
|
|
print("error: the following arguments are required: track", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
try:
|
|
separator = Separator(
|
|
model="htdemucs",
|
|
device=device,
|
|
shifts=shifts,
|
|
split=split,
|
|
overlap=overlap,
|
|
progress=True,
|
|
jobs=jobs,
|
|
segment=segment)
|
|
except ModelLoadingError as error:
|
|
fatal(error.args[0])
|
|
|
|
if stem != "" and stem not in separator.model.sources:
|
|
fatal(
|
|
'error: stem "{stem}" is not in selected model. '
|
|
"STEM must be one of {sources}.".format(
|
|
stem=stem, sources=", ".join(separator.model.sources)
|
|
)
|
|
)
|
|
out = out
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
print(f"Separated tracks will be stored in {out.resolve()}")
|
|
|
|
def random_name(length=8):
|
|
return ''.join(random.choices(string.ascii_lowercase + string.digits, k=length))
|
|
|
|
track["name"] = random_name()
|
|
|
|
origin, res = separator.separate_audio_file(track)
|
|
|
|
if audio_format == "mp3":
|
|
ext = "mp3"
|
|
elif audio_format == "flac":
|
|
ext = "flac"
|
|
else:
|
|
ext = "wav"
|
|
kwargs = {
|
|
"samplerate": separator.samplerate,
|
|
"bitrate": mp3_bitrate,
|
|
"preset": mp3_preset,
|
|
"clip": clip_mode,
|
|
"as_float": bit_depth == "float32",
|
|
"bits_per_sample": 24 if bit_depth == "int24" else 16,
|
|
}
|
|
|
|
outputs = {}
|
|
|
|
if not stem or stem.strip() == "":
|
|
for name, source in res.items():
|
|
output_path = out / filename.format(
|
|
track=track["name"].rsplit(".", 1)[0],
|
|
stem=name,
|
|
ext=ext,
|
|
)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
outputs[name] = save_audio(source, str(output_path), **kwargs)
|
|
|
|
else:
|
|
selected_stem = stem
|
|
minus_stem_path = out / filename.format(
|
|
track=track["name"].rsplit(".", 1)[0],
|
|
stem=f"minus_{selected_stem}",
|
|
ext=ext,
|
|
)
|
|
if other_method == "minus":
|
|
minus_stem_path.parent.mkdir(parents=True, exist_ok=True)
|
|
outputs["derived_stem"] = save_audio(origin - res[selected_stem], str(minus_stem_path), **kwargs)
|
|
stem_path = out / filename.format(
|
|
track=track["name"].rsplit(".", 1)[0],
|
|
stem=selected_stem,
|
|
ext=ext,
|
|
)
|
|
stem_path.parent.mkdir(parents=True, exist_ok=True)
|
|
outputs[selected_stem] = save_audio(res.pop(selected_stem), str(stem_path), **kwargs)
|
|
if other_method == "add":
|
|
other_stem = th.zeros_like(next(iter(res.values())))
|
|
for src in res.values():
|
|
other_stem += src
|
|
no_stem_path = out / filename.format(
|
|
track=track["name"].rsplit(".", 1)[0],
|
|
stem=f"no_{selected_stem}",
|
|
ext=ext,
|
|
)
|
|
no_stem_path.parent.mkdir(parents=True, exist_ok=True)
|
|
outputs["derived_stem"] = save_audio(other_stem, str(no_stem_path), **kwargs)
|
|
|
|
for source in self.OUTPUT_SOURCES:
|
|
if source not in outputs:
|
|
outputs[source] = {
|
|
"waveform": th.zeros(2, 1).unsqueeze(0),
|
|
"sample_rate": 44100
|
|
}
|
|
|
|
ordered_outputs = [outputs[source] for source in self.OUTPUT_SOURCES]
|
|
|
|
print(outputs)
|
|
drums, bass, others, vocals, derived_stem = ordered_outputs
|
|
|
|
return (drums, bass, others, vocals, derived_stem)
|