""" ╔════════════════════════════════════════╗ ║ ║ ║ 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)