203 lines
7.8 KiB
Python
203 lines
7.8 KiB
Python
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
|
||
|
||
print("Python executable:", sys.executable)
|
||
print("Python sys.path:", sys.path)
|
||
|
||
class DeepExtractV2Node:
|
||
"""
|
||
ComfyUI Node for Demucs-style separation.
|
||
Produces separate outputs for: drums, bass, other, vocals.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
# Ana girdiler
|
||
"track": ("AUDIO", {"default": None, "multiline": True, "tooltip": "Input audio track to separate"}),
|
||
|
||
# Çıktı klasörü
|
||
"out": ("STRING", {"default": "separated", "tooltip": "Folder where to put extracted track"}),
|
||
|
||
# Hesaplama cihazı
|
||
"device": (["cuda", "cpu"], {"default": "cuda", "tooltip": "Device to use for inference"}),
|
||
|
||
# Overlap değeri
|
||
"overlap": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "Overlap between audio chunks"}),
|
||
|
||
# Random shift sayısı
|
||
"shifts": ("INT", {"default": 1, "min": 1, "max": 20, "tooltip": "Number of random shifts for stabilization"}),
|
||
|
||
# Segment boyutu
|
||
"segment": ("INT", {"default": None, "min": 0, "max": 600, "tooltip": "Split size of each chunk (0 = no split)"}),
|
||
|
||
# Ses derinliği formatı
|
||
"bit_depth": (["default", "int24", "float32"], {"default": "default", "tooltip": "Audio output bit depth"}),
|
||
|
||
# Clip mode
|
||
"clip_mode": (["rescale", "clamp", "none"], {"default": "rescale", "tooltip": "Clipping handling strategy"}),
|
||
|
||
# Ses formatı
|
||
"audio_format": (["wav", "flac", "mp3"], {"default": "wav", "tooltip": "Output audio file format"}),
|
||
|
||
# MP3 bitrate
|
||
"mp3_bitrate": ("INT", {"default": 320, "min": 64, "max": 320, "tooltip": "Bitrate for MP3 output"}),
|
||
|
||
# MP3 encoder preset
|
||
"mp3_preset": ("INT", {"default": 2, "min": 2, "max": 7, "tooltip": "Encoder preset (2=best, 7=fast)"}),
|
||
|
||
# Paralel iş sayısı
|
||
"jobs": ("INT", {"default": 0, "min": 0, "max": 32, "tooltip": "Number of parallel jobs"}),
|
||
|
||
# Verbose modu
|
||
"verbose": ("BOOLEAN", {"default": False, "tooltip": "Enable verbose logging"}),
|
||
|
||
# Split aktif/pasif
|
||
"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}."}),
|
||
|
||
# Çıktı dosya adı formatı
|
||
"filename": ("STRING", {"default": "{track}/{stem}.{ext}", "tooltip": "Output filename format"}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = tuple(["AUDIO"] * len(SOURCES))
|
||
RETURN_NAMES = tuple(SOURCES)
|
||
FUNCTION = "separate"
|
||
CATEGORY = "Audio/Separation"
|
||
|
||
def separate(self, track, out, device, overlap, shifts,segment,bit_depth,clip_mode,audio_format,mp3_bitrate,mp3_preset,jobs,verbose,split,other_method,stem,filename):
|
||
print("🎧 Input values:")
|
||
print(f"tracks: {track}")
|
||
print(f"out: {out}")
|
||
print(f"device: {device}")
|
||
print(f"overlap: {overlap}")
|
||
print(f"shifts: {shifts}")
|
||
print(f"verbose: {verbose}")
|
||
|
||
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 is not "" 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()
|
||
|
||
print(f"Separating track {track}")
|
||
|
||
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 stem == "":
|
||
for name, source in res.items():
|
||
stem = out / filename.format(
|
||
track=track["name"].rsplit(".", 1)[0],
|
||
trackext=track["name"].rsplit(".", 1)[-1],
|
||
stem=name,
|
||
ext=ext,
|
||
)
|
||
stem.parent.mkdir(parents=True, exist_ok=True)
|
||
outputs[name] = save_audio(source, str(stem), **kwargs)
|
||
else:
|
||
stem = out / filename.format(
|
||
track=track["name"].rsplit(".", 1)[0],
|
||
trackext=track["name"].rsplit(".", 1)[-1],
|
||
stem="minus_" + stem,
|
||
ext=ext,
|
||
)
|
||
if other_method == "minus":
|
||
stem.parent.mkdir(parents=True, exist_ok=True)
|
||
outputs[name] = save_audio(origin - res[stem], str(stem), **kwargs)
|
||
stem = out / filename.format(
|
||
track=track["name"].rsplit(".", 1)[0],
|
||
trackext=track["name"].rsplit(".", 1)[-1],
|
||
stem=stem,
|
||
ext=ext,
|
||
)
|
||
stem.parent.mkdir(parents=True, exist_ok=True)
|
||
outputs[name] = save_audio(res.pop(stem), str(stem), **kwargs)
|
||
# Warning : after poping the stem, selected stem is no longer in the dict 'res'
|
||
if other_method == "add":
|
||
other_stem = th.zeros_like(next(iter(res.values())))
|
||
for i in res.values():
|
||
other_stem += i
|
||
stem = out / filename.format(
|
||
track=track["name"].rsplit(".", 1)[0],
|
||
trackext=track["name"].rsplit(".", 1)[-1],
|
||
stem="no_" + stem,
|
||
ext=ext,
|
||
)
|
||
stem.parent.mkdir(parents=True, exist_ok=True)
|
||
outputs[name] = save_audio(other_stem, str(stem), **kwargs)
|
||
|
||
drums, bass, others, vocals = outputs.values()
|
||
|
||
print(drums)
|
||
print(bass)
|
||
print(others)
|
||
print(vocals)
|
||
|
||
print(outputs)
|
||
|
||
return (drums, bass, others, vocals) |